Problem Description

You are in a mystic dungeon with n magicians standing in a line. Each magician has an energy value that can be positive (giving you energy) or negative (taking energy from you).

You've been cursed with a teleportation spell: after absorbing energy from magician at position i, you automatically teleport to magician at position (i + k). This continues until you reach a position where (i + k) doesn't exist (goes beyond the array).

Your journey works as follows:

• Choose any starting magician position
• Absorb that magician's energy (you must take it, whether positive or negative)
• Teleport forward by exactly k positions
• Repeat until you can't teleport anymore (would go out of bounds)

Given an array energy representing each magician's energy value and an integer k representing the teleport distance, find the maximum total energy you can gain by choosing the optimal starting position.

For example, if energy = [5, -3, 2, -1, 4] and k = 2:

• Starting at index 0: you'd visit positions 0 → 2 → 4, gaining 5 + 2 + 4 = 11 energy
• Starting at index 1: you'd visit positions 1 → 3, gaining -3 + (-1) = -4 energy
• Starting at index 2: you'd visit position 2 only, gaining 2 energy
• Starting at index 3: you'd visit position 3 only, gaining -1 energy
• Starting at index 4: you'd visit position 4 only, gaining 4 energy

The maximum energy would be 11 by starting at index 0.

Intuition

The key insight is recognizing that magicians form groups based on the jump distance k. When we jump by k positions each time, we're essentially visiting every k-th element starting from our initial position.

Think about it this way: if we start at position i, we visit positions i, i+k, i+2k, i+3k, and so on until we go out of bounds. This means the array is divided into k independent chains, where:

• Chain 0 contains positions: 0, k, 2k, 3k, ...
• Chain 1 contains positions: 1, k+1, 2k+1, 3k+1, ...
• Chain 2 contains positions: 2, k+2, 2k+2, 3k+2, ...
• And so on...

Each starting position belongs to exactly one chain, and once we start on a chain, we must follow it to the end. The crucial observation is that we want to maximize our energy gain, which means we might not want to start at the beginning of a chain - we might want to skip some negative values at the start and begin somewhere in the middle.

Since we must move forward by k each time, the last possible positions we can visit are in the range [n-k, n-1]. These are the potential ending positions of our journey. Working backwards from these endpoints, we can calculate the sum of energy values at each k interval.

By iterating through each possible endpoint and traversing backwards (subtracting k each time), we can compute all possible suffix sums for each chain. As we build these suffix sums, we keep track of the maximum sum encountered. This approach ensures we find the optimal starting position that gives us the maximum energy, as it considers all possible sub-paths within each chain.

The backwards traversal is elegant because it naturally handles the constraint that we must take all energy values from our starting point to the end - we're essentially asking "what's the best suffix of each chain to take?"

Learn more about Prefix Sum patterns.

Solution Approach

The implementation uses the enumeration and suffix sum strategy to find the maximum energy gain.

We start by initializing ans = -inf to track the maximum energy found so far, and get the array length n.

The core of the algorithm is a loop that iterates through all possible ending positions:

for i in range(n - k, n):

This range [n-k, n) represents the last k positions in the array. These are the only positions that can serve as endpoints of our journey since any position before n-k would have at least one more jump of size k that stays within bounds.

For each ending position i, we traverse backwards through its chain:

j, s = i, 0
while j >= 0:
    s += energy[j]
    ans = max(ans, s)
    j -= k

Here's what happens in this inner loop:

• j starts at the ending position i and moves backwards by k each iteration
• s accumulates the energy sum as we move backwards through the chain
• After adding each energy value to s, we update ans with the maximum between current ans and s

The key insight is that by building the sum incrementally while moving backwards, we're calculating all possible suffix sums for this chain. For example, if we're traversing positions [2, 7, 12] backwards (starting from position 12), we calculate:

• Sum from position 12 only: energy[12]
• Sum from positions 7 to 12: energy[12] + energy[7]
• Sum from positions 2 to 12: energy[12] + energy[7] + energy[2]

By updating ans after each addition, we capture the maximum energy achievable by starting at any position within this chain.

The algorithm effectively checks all k chains (one for each ending position in [n-k, n)), and for each chain, it evaluates all possible starting positions by computing suffix sums. The time complexity is O(n) since each element is visited at most once across all chains, and the space complexity is O(1) as we only use a few variables to track the current state.

Example Walkthrough

Let's walk through a concrete example with energy = [3, -2, 5, -1, 4, -3, 2] and k = 3.

First, we identify our possible ending positions. Since n = 7 and k = 3, the ending positions are in range [n-k, n) = [4, 7), which means indices 4, 5, and 6.

Processing ending position 4 (index 4):

• Start at position 4: s = 0 + energy[4] = 0 + 4 = 4, update ans = 4
• Jump back to position 1: s = 4 + energy[1] = 4 + (-2) = 2, update ans = max(4, 2) = 4
• Jump back to position -2: Out of bounds, stop

Chain visited: [1, 4] with suffix sums: [4, 2]

Processing ending position 5 (index 5):

• Start at position 5: s = 0 + energy[5] = 0 + (-3) = -3, update ans = max(4, -3) = 4
• Jump back to position 2: s = -3 + energy[2] = -3 + 5 = 2, update ans = max(4, 2) = 4
• Jump back to position -1: Out of bounds, stop

Chain visited: [2, 5] with suffix sums: [-3, 2]

Processing ending position 6 (index 6):

• Start at position 6: s = 0 + energy[6] = 0 + 2 = 2, update ans = max(4, 2) = 4
• Jump back to position 3: s = 2 + energy[3] = 2 + (-1) = 1, update ans = max(4, 1) = 4
• Jump back to position 0: s = 1 + energy[0] = 1 + 3 = 4, update ans = max(4, 4) = 4
• Jump back to position -3: Out of bounds, stop

Chain visited: [0, 3, 6] with suffix sums: [2, 1, 4]

The maximum energy we can gain is 4, which can be achieved by:

• Starting at index 4 and taking only that magician's energy (4)
• OR starting at index 0 and visiting positions 0→3→6, gaining 3 + (-1) + 2 = 4

The algorithm correctly identifies both optimal paths by computing all possible suffix sums across the chains.

Solution Implementation

Time and Space Complexity

The time complexity is O(n), where n is the length of the array energy.

The algorithm iterates through k starting positions (from index n-k to n-1). For each starting position i, it traverses backwards with step size k until reaching the beginning of the array. The key insight is that each element in the array is visited at most once across all iterations. This is because each element belongs to exactly one "chain" when jumping backwards by k positions, and we only process each chain once by starting from its last possible position in the range [n-k, n-1]. Therefore, the total number of operations is proportional to n.

The space complexity is O(1). The algorithm only uses a constant amount of extra space for variables ans, i, j, and s, regardless of the input size.

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Misunderstanding the Direction of Traversal

Pitfall: A common mistake is thinking we need to traverse forward from each starting position to find the maximum energy. This leads to inefficient code that checks every possible starting position:

# Inefficient approach - checking all starting positions
def maximumEnergy(energy: List[int], k: int) -> int:
    max_energy = -inf
    n = len(energy)
  
    # Wrong: Trying every position as a starting point
    for start in range(n):
        current_sum = 0
        pos = start
        while pos < n:
            current_sum += energy[pos]
            pos += k
        max_energy = max(max_energy, current_sum)
  
    return max_energy

Why it's wrong: This approach has O(n²/k) time complexity because we're checking all n starting positions, when we only need to check k distinct chains.

Solution: Recognize that positions form k independent chains based on their modulo k value. By starting from the end positions and working backwards, we can efficiently compute all possible sums in O(n) time.

2. Not Updating Maximum After Each Addition

Pitfall: Only checking the total sum at the end of each chain traversal:

# Wrong: Only checking complete chain sums
def maximumEnergy(energy: List[int], k: int) -> int:
    max_energy = -inf
    n = len(energy)
  
    for i in range(n - k, n):
        j, s = i, 0
        while j >= 0:
            s += energy[j]
            j -= k
        # Wrong: Only updating after the complete chain
        max_energy = max(max_energy, s)
  
    return max_energy

Why it's wrong: This misses optimal solutions where starting from the middle of a chain gives better results. For example, with energy = [-5, 10, -3] and k = 2, starting at index 2 gives -3 (better than starting at index 0 which gives -5 + (-3) = -8).

Solution: Update the maximum after each energy addition to capture all possible suffix sums:

while j >= 0:
    s += energy[j]
    max_energy = max(max_energy, s)  # Update after EACH addition
    j -= k

3. Incorrect Range for Ending Positions

Pitfall: Using the wrong range for potential ending positions:

# Wrong: Incorrect range
for i in range(n - k + 1, n + 1):  # Off by one error
    # ...

or

# Wrong: Starting from n - k - 1
for i in range(n - k - 1, n):  # Includes positions that aren't endpoints
    # ...

Why it's wrong: The correct ending positions are exactly the last k positions: indices [n-k, n-k+1, ..., n-1]. Any position before n-k would have another position k steps ahead within the array bounds.

Solution: Use the precise range range(n - k, n) which gives us exactly the last k positions of the array.

4. Forgetting Edge Cases

Pitfall: Not handling cases where k equals or exceeds n:

# Potential issue when k >= n
for i in range(n - k, n):  # Could result in negative start of range
    # ...

Why it's problematic: When k >= n, every position becomes its own independent chain (no jumps possible), and range(n - k, n) might not cover all positions correctly.

Solution: Add validation or ensure the algorithm handles this naturally:

def maximumEnergy(energy: List[int], k: int) -> int:
    max_energy = -inf
    n = len(energy)
  
    # When k >= n, each position is its own chain
    start_range = max(0, n - k)
  
    for i in range(start_range, n):
        j, s = i, 0
        while j >= 0:
            s += energy[j]
            max_energy = max(max_energy, s)
            j -= k
  
    return max_energy