Problem Statement in English

You’re given a 0-indexed integer array power representing the power of different spells. You can cast any spell only once.

If you cast a spell with power x, you will deal x damage and all spells with power x - 1, x - 2, x + 1 and x + 2 will become unavailable to cast.

Return the maximum total damage you can deal by casting the spells optimally.

Approach

We can extract the unique powers from the power array and sort them. We also need a mapping from each power to its occurences so that we can easily use any further spells of the same power again.

We also need to skip the next and potentially the next to next power if we decide to use a spell of a certain power. This can be efficiently done by checking the mapping to see if the next power and next to next power exists or not.

We can then use dynamic programming to decide whether to include or exclude each unique power.

Solution in Python


class Solution:
    def maximumTotalDamage(self, power: List[int]) -> int:
        c = Counter(power)
        power = sorted(set(power))
        res = 0
        l = len(power)

        @cache
        def dp(i):
            if i >= l:
                return 0

            # exclude
            exclude = dp(i + 1)

            # include
            step = 1
            if power[i] + 1 in c:
                step += 1
            if power[i] + 2 in c:
                step += 1
            include = power[i] * c[power[i]] + dp(i + step)

            return max(exclude, include)

        for i in range(l):
            res = max(res, dp(i))

        return res

Complexity

  • Time: $O(n \log n)$, where $n$ is the length of the power array. This is due to the sorting step and the dynamic programming approach.

  • Space: $O(n)$, for storing the count of each unique power and the memoization cache.

And we are done.