2172. Maximum AND Sum of Array

[zeikar]

Problem link

https://leetcode.com/problems/maximum-and-sum-of-array/

Problem Summary

슬롯이 있고 숫자를 한 슬롯에 최대 2개까지 집어넣을 수 있다. 이때 슬롯의 번호와 숫자를 AND 연산한 합 중에 최대를 구하는 문제.

Solution

비트마스크 DP 이다.
처음에는 슬롯에 2개씩 들어가므로 슬롯의 크기를 2배로 해서 돌려봤는데 시간 초과... 결국 비트마스크를 2개 써서 통과.

dp[i][slot1][slot2] = slot1, slot2에 숫자가 있을 때의 최대 합.

Source Code

from functools import lru_cache
from typing import List


class Solution:
    def maximumANDSum(self, nums: List[int], numSlots: int) -> int:
        @lru_cache(None)
        def solve(idx, slot1, slot2):
            if idx == len(nums):
                return 0

            ret = 0
            for slot in range(numSlots):
                if slot1 & (1 << slot) == 0:
                    ret = max(ret, solve(idx + 1, slot1 | (1 << slot), slot2) + (nums[idx] & (slot + 1)))
                elif slot2 & (1 << slot) == 0:
                    ret = max(ret, solve(idx + 1, slot1, slot2 | (1 << slot)) + (nums[idx] & (slot + 1)))

            return ret

        return solve(0, 0, 0)