Problem Description

You are given an array power where each element represents the damage value of a spell that a magician can cast. Multiple spells can have the same damage value.

The key constraint is: if the magician casts a spell with damage power[i], they cannot cast any spell with damage values of power[i] - 2, power[i] - 1, power[i] + 1, or power[i] + 2. In other words, casting a spell blocks out a range of 5 consecutive damage values (the spell itself and 2 values on each side).

Each spell can only be cast once, meaning you can only use each element in the array once.

Your task is to find the maximum total damage the magician can achieve by selecting spells to cast while respecting the blocking constraint.

For example, if power = [1, 1, 3, 4], the magician could:

• Cast both spells with damage 1 (total damage = 2), which would block damage values -1, 0, 1, 2, 3
• Or cast the spell with damage 4 (total damage = 4), which would block damage values 2, 3, 4, 5, 6

The maximum would be 4 in this case.

Intuition

The problem resembles the classic "House Robber" problem where we need to make choices about which elements to take while avoiding certain constraints. Here, the constraint is that selecting a damage value blocks a range of 5 consecutive values.

First, let's think about the structure of the problem. If we sort the array, we can make decisions in order - for each unique damage value, we either take all occurrences of it or skip it entirely. This is because if we can take one spell with damage x, we can take all spells with damage x since they don't interfere with each other.

The key insight is that when we decide to use all spells of damage value x, we must skip to the next "safe" damage value - one that is greater than x + 2. This is where binary search becomes useful: after sorting, we can efficiently find the index of the first damage value that's greater than x + 2.

Now we can frame this as a dynamic programming problem. For each position i in the sorted array, we have two choices:

1. Skip the current damage value: Move to the next different damage value (skip all occurrences of power[i])
2. Take the current damage value: Collect power[i] * count[power[i]] damage and jump to the first safe position

The challenge is handling duplicates efficiently. By using a counter to track occurrences, we can process all duplicates at once. When we skip, we skip all occurrences; when we take, we take all occurrences.

This leads to a recursive solution with memoization: at each position, we compute the maximum damage we can get from that position onwards by comparing the two choices above. The overlapping subproblems (same positions being evaluated multiple times) make memoization effective.

Learn more about Two Pointers, Binary Search, Dynamic Programming and Sorting patterns.

Solution Approach

The implementation follows these key steps:

1. Preprocessing the Data

First, we sort the power array and create a counter cnt to track how many times each damage value appears:

cnt = Counter(power)
power.sort()

2. Building the Next Valid Index Array

For each position i, we precompute where we can jump to if we decide to use the damage value at position i. This is the index of the first element greater than power[i] + 2:

nxt = [bisect_right(power, x + 2, lo=i + 1) for i, x in enumerate(power)]

The bisect_right function performs binary search to find this position efficiently. We start searching from lo=i + 1 to optimize the search range.

3. Dynamic Programming with Memoization

We define a recursive function dfs(i) that calculates the maximum damage starting from index i:

@cache
def dfs(i: int) -> int:
    if i >= n:
        return 0
    a = dfs(i + cnt[power[i]])
    b = power[i] * cnt[power[i]] + dfs(nxt[i])
    return max(a, b)

The function works as follows:

• Base case: If i >= n, we've processed all elements, return 0
• Choice 1 (Skip): Skip all occurrences of power[i] by jumping to i + cnt[power[i]]. This gives us damage value a = dfs(i + cnt[power[i]])
• Choice 2 (Take): Take all occurrences of power[i] (gaining power[i] * cnt[power[i]] damage) and jump to the next valid position nxt[i]. This gives us damage value b = power[i] * cnt[power[i]] + dfs(nxt[i])
• Return the maximum of both choices

The @cache decorator automatically memoizes the results, preventing redundant calculations for the same index.

4. Getting the Final Answer

The maximum total damage is obtained by calling dfs(0), which considers all possible combinations starting from the first element.

Time Complexity: O(n log n) for sorting and binary search operations, plus O(n) for the DP traversal, giving O(n log n) overall.

Space Complexity: O(n) for the counter, next array, and memoization cache.

Example Walkthrough

Let's walk through the solution with power = [1, 1, 3, 4, 7].

Step 1: Preprocessing

• Sort the array: power = [1, 1, 3, 4, 7] (already sorted)
• Create counter: cnt = {1: 2, 3: 1, 4: 1, 7: 1}
• Array length: n = 5

Step 2: Build Next Valid Index Array For each position, find the first index where the value is greater than power[i] + 2:

• i=0: power[0]=1, need value > 3. Binary search finds index 3 (value 4)
• i=1: power[1]=1, need value > 3. Binary search finds index 3 (value 4)
• i=2: power[2]=3, need value > 5. Binary search finds index 4 (value 7)
• i=3: power[3]=4, need value > 6. Binary search finds index 4 (value 7)
• i=4: power[4]=7, need value > 9. Binary search finds index 5 (out of bounds)

So nxt = [3, 3, 4, 4, 5]

Step 3: Dynamic Programming Traversal Starting from dfs(0):

dfs(0): Process damage value 1 (at index 0)

• Skip option: Jump to index 0 + cnt[1] = 0 + 2 = 2, call dfs(2)
• Take option: Gain 1 * 2 = 2 damage, jump to index nxt[0] = 3, call dfs(3)

dfs(2): Process damage value 3 (at index 2)

• Skip option: Jump to index 2 + cnt[3] = 2 + 1 = 3, call dfs(3)
• Take option: Gain 3 * 1 = 3 damage, jump to index nxt[2] = 4, call dfs(4)

dfs(3): Process damage value 4 (at index 3)

• Skip option: Jump to index 3 + cnt[4] = 3 + 1 = 4, call dfs(4)
• Take option: Gain 4 * 1 = 4 damage, jump to index nxt[3] = 4, call dfs(4)

dfs(4): Process damage value 7 (at index 4)

• Skip option: Jump to index 4 + cnt[7] = 4 + 1 = 5, call dfs(5)
• Take option: Gain 7 * 1 = 7 damage, jump to index nxt[4] = 5, call dfs(5)

dfs(5): Base case, return 0

Step 4: Backtracking with Values

• dfs(5) = 0
• dfs(4) = max(0, 7 + 0) = 7
• dfs(3) = max(dfs(4), 4 + dfs(4)) = max(7, 4 + 7) = 11
• dfs(2) = max(dfs(3), 3 + dfs(4)) = max(11, 3 + 7) = 11
• dfs(0) = max(dfs(2), 2 + dfs(3)) = max(11, 2 + 11) = 13

Result: Maximum damage = 13

This is achieved by taking both spells with damage 1 (total 2) and both spells with damage 4 and 7 (total 11), giving 2 + 4 + 7 = 13. The spell with damage 3 is blocked by taking damage 1.

Solution Implementation

Time and Space Complexity

Time Complexity: O(n log n)

The time complexity is dominated by:

• Sorting the power array: O(n log n)
• Building the nxt array using bisect_right for each element: O(n log n) (each binary search takes O(log n) and we do it n times)
• The DFS with memoization visits each unique index at most once due to the @cache decorator, resulting in O(n) calls
• Each DFS call performs O(1) operations (aside from recursive calls)

Therefore, the overall time complexity is O(n log n).

Space Complexity: O(n)

The space complexity consists of:

• cnt Counter dictionary: O(n) in worst case (all elements unique)
• Sorted power array: O(n)
• nxt array: O(n)
• DFS recursion stack depth: O(n) in worst case
• Memoization cache: O(n) for storing results of up to n different indices

The total space complexity is O(n).

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

Common Pitfalls

1. Not Handling Duplicate Values Correctly

The Pitfall: A common mistake is treating each element in the array independently without considering that multiple spells can have the same damage value. If you process each element one by one, you might incorrectly skip or take individual occurrences of the same value, leading to suboptimal solutions.

Why It Happens: The problem states "each spell can only be cast once," which might lead to thinking we need to make decisions for each array element individually. However, when we decide to cast a spell with a certain damage value, we should cast ALL available spells with that exact damage value to maximize damage.

Example of Incorrect Approach:

# WRONG: Processing each element individually
def dfs(i):
    if i >= n:
        return 0
    # Skip current element
    skip = dfs(i + 1)
    # Take current element and jump to next valid
    take = power[i] + dfs(next_valid[i])
    return max(skip, take)

The Solution: Group identical damage values together and make decisions for the entire group. When we choose to use a damage value, we use ALL occurrences of it:

# CORRECT: Process all duplicates together
def dfs(i):
    if i >= n:
        return 0
    # Skip ALL occurrences of power[i]
    skip = dfs(i + cnt[power[i]])
    # Take ALL occurrences of power[i]
    take = power[i] * cnt[power[i]] + dfs(nxt[i])
    return max(skip, take)

2. Incorrect Next Valid Index Calculation

The Pitfall: Miscalculating where to jump after selecting a damage value. Common errors include:

• Using bisect_left instead of bisect_right
• Looking for values greater than power[i] + 1 instead of power[i] + 2
• Not optimizing the search range with lo=i+1

Why It Happens: The constraint blocks values in the range [power[i] - 2, power[i] + 2]. The next valid value must be strictly greater than power[i] + 2, not greater than or equal to.

Example of Incorrect Approaches:

# WRONG: Using bisect_left (might include power[i] + 2)
nxt = [bisect_left(power, x + 3) for x in power]

# WRONG: Looking for wrong boundary
nxt = [bisect_right(power, x + 1) for x in power]

# INEFFICIENT: Not optimizing search range
nxt = [bisect_right(power, x + 2) for x in power]

The Solution: Use bisect_right to find the first index with value strictly greater than power[i] + 2:

# CORRECT: Find first value > power[i] + 2, starting from i+1
nxt = [bisect_right(power, x + 2, lo=i + 1) for i, x in enumerate(power)]

3. Not Sorting the Array First

The Pitfall: Attempting to solve the problem without sorting leads to incorrect results because:

• The binary search for next valid index won't work correctly
• The DP state transitions become invalid when jumping indices

Why It Happens: The original problem doesn't explicitly mention sorting, and one might think the order of spells matters.

The Solution: Always sort the array first. The order of casting spells doesn't affect the total damage - only which spells are cast matters. Sorting enables efficient binary search and proper DP transitions.

power.sort()  # Essential first step