<p>A magician has various spells.</p>

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

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

<p>Each spell can be cast <strong>only once</strong>.</p>

<p>Return the <strong>maximum</strong> possible <em>total damage</em> that a magician can cast.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>

<div class="example-block">
<p><strong>Input:</strong> <span class="example-io">power = [1,1,3,4]</span></p>

<p><strong>Output:</strong> <span class="example-io">6</span></p>

<p><strong>Explanation:</strong></p>

<p>The maximum possible damage of 6 is produced by casting spells 0, 1, 3 with damage 1, 1, 4.</p>
</div>

<p><strong class="example">Example 2:</strong></p>

<div class="example-block">
<p><strong>Input:</strong> <span class="example-io">power = [7,1,6,6]</span></p>

<p><strong>Output:</strong> <span class="example-io">13</span></p>

<p><strong>Explanation:</strong></p>

<p>The maximum possible damage of 13 is produced by casting spells 1, 2, 3 with damage 1, 6, 6.</p>
</div>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>1 &lt;= power.length &lt;= 10<sup>5</sup></code></li>
	<li><code>1 &lt;= power[i] &lt;= 10<sup>9</sup></code></li>
</ul>

class Solution:
    def maximumTotalDamage(self, power: List[int]) -> int:
        '''
        At idx i, see if you want to pick it or no pick it
        return max(pick, no_pick)
        How to correctly pick? 
        You can pick duplicates, so add all dups
        when you reach end of dup element and then call i+1
        If I'm at idx i, i should ensure that I don't pick power[i]-2/-1/+1/+2
        when i reach an idx, i'll do a look up for +1/+2
        '''
        n = len(power)
        power.sort() # Order doesn't matter, can sort and just look ahead
        memo = [0] * len(power)
        # LOOK AT THE PROBLEM CONSTRAINTS
        # 10^5 X 10^9 is 10^14 -- you can't cache that with two params, it's
        # too large of an array to store states
        # Memorize all the states you calculate 
        # If you want to see the optimization, see the last failed solution
        # which was building a massive 2D array in memory
        # By memoizing, we're solving this in 10^5 space max

        def get_next_best_index(curr_idx, spell):
            # Since array is sorted, we can binary search it
            low, hi = curr_idx, n-1
            while low<=hi:
                mid = (low+hi)//2
                if power[mid]==spell+1 or power[mid]==spell+2:
                    low = mid + 1
                else:
                    hi = mid - 1
            return low

        def idx_after_duplicate(curr_idx, spell):
            # Get the last index of the dup - i.e. rightmost boundary
            lo, hi = curr_idx, n - 1
            while lo <= hi:
                mid = (lo + hi)//2
                if power[mid]==spell:
                    lo = mid + 1
                else:
                    hi = mid - 1
            return lo

        def calc_damage(idx):
            if idx == len(power): 
                return 0
            if memo[idx]:
                return memo[idx]

            # Should I pick this spell? Let's find out
            pick = 0

            # Handle duplicates
            next_idx = idx_after_duplicate(idx, power[idx])
            tot_damage = power[idx] * (next_idx-idx)

            # Find the next index. Since the array is sorted, we can skip to 
            # some (last/final)index 'i+1' where power[i]==power[idx] + (1 or 2)
            # because if we land on any of these indices then we can't consider
            # the damage (that's the problem constraint)
            next_best_index = get_next_best_index(next_idx, power[idx])
            pick = calc_damage(next_best_index) + tot_damage

            # Or should I see if there's something better
            # If I don't pick this idx ?
            no_pick = calc_damage(next_idx)
            memo[idx] = max(pick, no_pick)
            return memo[idx]

        res = calc_damage(0)
        return res

71315