LeetCode 3147 - Taking Maximum Energy From the Mystic Dungeon

The problem asks us to maximize the total energy gained from a sequence of magicians arranged in a line, where each magician provides a certain energy value, which can be negative or positive.

leetcodemediumarraydynamic-programmingprefix-sum

LeetCode Problem 3147

Difficulty: 🟡 Medium
Topics: Array, Dynamic Programming, Prefix Sum

Solution

Problem Understanding

The problem asks us to maximize the total energy gained from a sequence of magicians arranged in a line, where each magician provides a certain energy value, which can be negative or positive. The catch is that after taking energy from a magician at index i, you are instantly teleported to magician (i + k) and continue this process until you go beyond the bounds of the array. The goal is to determine the starting index that produces the maximum total energy collected.

The input consists of an array energy of length n representing the energy from each magician and an integer k representing the teleport distance. The output is a single integer: the maximum energy that can be obtained following the teleport rule.

Constraints are significant: n can be up to 100,000, meaning a naive approach iterating through all starting points and summing energy for each possible teleport sequence could be too slow. Energy values can be negative, so simply summing largest elements will not work; the order and jumps matter.

Edge cases include all negative energy, k equal to 1 (every magician visited), and k equal to n-1 (only one jump possible).

Approaches

Brute Force: The simplest approach is to try starting at every index from 0 to n-1, simulate the teleport sequence by adding energy at each step (i, i+k, i+2k, ...), and track the maximum total. This guarantees correctness but will be too slow for large n because it requires O(n^2/k) operations in the worst case.

Optimal Approach: The key insight is that the problem can be solved with dynamic programming from right to left. Let dp[i] represent the maximum energy obtainable starting at magician i. The recurrence relation is dp[i] = energy[i] + (dp[i+k] if i+k < n else 0). By iterating from the end of the array backward, we can compute each dp[i] in O(1) time using previously computed results. The maximum of all dp[i] values gives the answer. This approach reduces the time complexity to O(n) and uses O(n) additional space for the DP array, which can be optimized to O(k) using a rolling array if desired.

Approach Time Complexity Space Complexity Notes
Brute Force O(n^2/k) O(1) Iterate all start indices and simulate each sequence
Optimal O(n) O(n) DP from end to start using recurrence dp[i] = energy[i] + dp[i+k]

Algorithm Walkthrough

  1. Initialize a DP array of length n to store the maximum energy starting from each magician.
  2. Iterate backward from the last magician to the first.
  3. For magician i, calculate dp[i] = energy[i] + (dp[i+k] if i+k < n else 0). This ensures we capture all energy from the current sequence starting at i.
  4. Keep track of the maximum value encountered during this iteration; this is the maximum energy obtainable.
  5. Return the maximum value.

Why it works: The recurrence captures the exact sequence of energy collection defined by the teleport rule. Since each dp[i] stores the maximum energy from starting at i, taking the maximum of all dp[i] ensures we find the optimal starting index. The right-to-left traversal guarantees that when calculating dp[i], the required dp[i+k] value is already computed.

Python Solution

from typing import List

class Solution:
    def maximumEnergy(self, energy: List[int], k: int) -> int:
        n = len(energy)
        dp = [0] * n
        max_energy = float('-inf')
        
        for i in range(n-1, -1, -1):
            if i + k < n:
                dp[i] = energy[i] + dp[i + k]
            else:
                dp[i] = energy[i]
            max_energy = max(max_energy, dp[i])
        
        return max_energy

Implementation Explanation: The DP array dp holds maximum energy from each index. We iterate from right to left to ensure we have already computed dp[i+k] before using it. At each step, we either add dp[i+k] to the current energy if within bounds, or just take the current energy. The max_energy variable keeps track of the best starting point seen so far.

Go Solution

func maximumEnergy(energy []int, k int) int {
    n := len(energy)
    dp := make([]int, n)
    maxEnergy := -1 << 31 // minimum integer
    
    for i := n-1; i >= 0; i-- {
        if i+k < n {
            dp[i] = energy[i] + dp[i+k]
        } else {
            dp[i] = energy[i]
        }
        if dp[i] > maxEnergy {
            maxEnergy = dp[i]
        }
    }
    
    return maxEnergy
}

Go-specific Notes: Go does not have inf for integers, so maxEnergy is initialized to the minimum possible integer. Slices are used for DP array. Logic mirrors Python version exactly.

Worked Examples

Example 1: energy = [5, 2, -10, -5, 1], k = 3

i dp[i] calculation dp[i] value
4 i+k >= n → dp[4] = 1 1
3 i+k >= n → dp[3] = -5 -5
2 i+k >= n → dp[2] = -10 -10
1 1+3=4 → dp[1] = 2 + dp[4] = 2+1 3
0 0+3=3 → dp[0] = 5 + dp[3] = 5-5 0

Max energy = 3

Example 2: energy = [-2, -3, -1], k = 2

i dp[i] calculation dp[i] value
2 i+k >= n → dp[2] = -1 -1
1 i+k >= n → dp[1] = -3 -3
0 0+2=2 → dp[0] = -2 + dp[2] = -2-1 -3

Max energy = -1

Complexity Analysis

Measure Complexity Explanation
Time O(n) We iterate once through the energy array calculating dp values
Space O(n) DP array stores maximum energy for each index

We could optimize space to O(k) by only storing the last k dp values if needed.

Test Cases

# Provided examples
assert Solution().maximumEnergy([5,2,-10,-5,1], 3) == 3  # Example 1
assert Solution().maximumEnergy([-2,-3,-1], 2) == -1      # Example 2

# Edge cases
assert Solution().maximumEnergy([1], 1) == 1             # Single element
assert Solution().maximumEnergy([-1], 1) == -1           # Single negative element
assert Solution().maximumEnergy([1000]*1000, 500) == 2000  # Large positive elements
assert Solution().maximumEnergy([-1000]*1000, 500) == -1000 # Large negative elements
assert Solution().maximumEnergy([1,-2,3,-4,5,-6], 2) == 5 # Mixed signs, optimal not starting at index 0
Test Why
[5,2,-10,-5,1], k=3 Tests mixed positive/negative with jumps
[-2,-3,-1], k=2 All negative numbers
[1], k=1 Minimum size input
[-1], k=1 Minimum size negative input
[1000]*1000, k=500 Large input, all positive
[-1000]*1000, k=500 Large input, all negative
[1,-2,3,-4,5,-6], k=2 Mixed signs, optimal starting index is not 0

Edge Cases

All negative energies: If all magicians have negative energy, the algorithm correctly selects the least negative starting point. The DP recurrence ensures that any further jumps only reduce total energy, so the max is correctly identified.

Single magician: For n=1, the loop handles the base case as i+k >= n triggers taking only the current energy. No special branching is needed.

k equals array length minus one: In this case, each magician only reaches one other possible magician. The algorithm correctly handles i+k >= n and returns either energy[i] or energy[i]+dp[i+k] if in bounds. No out-of-bounds access occurs.