Problem Link

Description


A magician has various spells.

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

It is a known fact that if a magician decides to cast a spell with a damage of power[i], they cannot cast any spell with a damage of power[i] - 2, power[i] - 1, power[i] + 1, or power[i] + 2.

Each spell can be cast only once.

Return the maximum possible total damage that a magician can cast.

 

Example 1:

Input: power = [1,1,3,4]

Output: 6

Explanation:

The maximum possible damage of 6 is produced by casting spells 0, 1, 3 with damage 1, 1, 4.

Example 2:

Input: power = [7,1,6,6]

Output: 13

Explanation:

The maximum possible damage of 13 is produced by casting spells 1, 2, 3 with damage 1, 6, 6.

 

Constraints:

  • 1 <= power.length <= 105
  • 1 <= power[i] <= 109

Solution


Python3

class Solution:
    def maximumTotalDamage(self, power: List[int]) -> int:
        counter = Counter(power)
        A = sorted(counter.keys())
        N = len(counter)
        
        @cache
        def go(index):
            if index == N:
                return 0
            
            # skip
            res = go(index + 1)
            
            # take
            nextIndex = index + 1
            while nextIndex < N and A[nextIndex] - A[index] <= 2:
                nextIndex += 1
            
            res = max(res, A[index] * counter[A[index]] + go(nextIndex))
            
            return res
        
        return go(0)