Maximum AND Sum of Array LeetCode solution in Python

1
# Runtime: 1881 ms (Top 9.44%) | Memory: 23.9 MB (Top 80.86%)
2
from functools import lru_cache, cache
3

4

5
class Solution:
6
    def maximumANDSum(self, nums: List[int], numSlots: int) -> int:
7
        @cache
8
        def dp(i=0, m1=0, m2=0):  # mask1, mask2
9
            if i == len(nums):
10
                return 0
11
            ans = 0
12
            for s in range(numSlots):
13
                if m2 & (1 << s) == 0:  # i.e. 0b0?, implying the slot is not full
14
                    if m1 & (1 << s) == 0:  # 0b00 + 1 => 0b01
15
                        nm1 = m1 | (1 << s)
16
                        nm2 = m2
17
                    else:  # 0b01 + 1 => 0b10
18
                        nm1 = m1 & ~(1 << s)
19
                        nm2 = m2 | (1 << s)
20
                    ans = max(
21
                        ans, dp(i + 1, nm1, nm2) + ((s + 1) & nums[i])
22
                    )  # s + 1 is the actual slot no.
23
            return ans
24

25
        return dp()