[LeetCode] 698. Partition to K Equal Sum Subsets (Medium)

이 문제를 풀기 위해서는 비트마스크 DP에 대해서 알아야 한다.

배열의 길이가 최대 16이기 때문에, 각 원소를 썼는지 안 썼는지를 하나의 비트마스크로 표현할 수 있다. 그리고 어떤 원소들을 이미 사용했는지 알 수 있으면, 지금 만들고 있는 부분집합의 합이 목표값에서 어디까지 차 있는지도 계산할 수 있다.

0. 필요한 개념 먼저 정리하기

먼저 이 문제의 목표는 k개의 부분집합을 만드는 것이다. 모든 부분집합의 합이 같아야 하므로, 전체 합을 k로 나눈 값이 각 부분집합의 목표 합이 된다.

예를 들어 nums = [4,3,2,3,5,2,1], k = 4라면 전체 합은 20이고 각 부분집합의 목표 합은 5다. 따라서 (5), (4,1), (3,2), (3,2)처럼 모든 그룹이 합 5로 끝나야 한다.

여기서 중요한 전제는 두 가지다.

• 전체 합이 k로 나누어떨어지지 않으면 답은 바로 false다.
• 어떤 숫자가 목표 합보다 크면 그 숫자는 어떤 부분집합에도 들어갈 수 없으므로 답은 false다.

그다음 필요한 개념이 비트마스크다. 배열의 원소가 최대 16개이므로, 2^16개의 사용 상태를 충분히 탐색할 수 있다. 예를 들어 정렬된 배열이 [5, 4, 3, 3, 2, 2, 1]일 때, 마스크 0000011은 앞의 두 원소 5, 4를 이미 사용했다는 뜻으로 볼 수 있다.

DP 상태는 다음처럼 정의한다.

dp[mask] = mask에 포함된 원소들을 사용했을 때,
           현재 채우는 중인 부분집합의 합

다만 현재 부분집합의 합은 항상 target보다 작게 유지하면 된다. 부분집합 하나가 정확히 target이 되는 순간에는 다음 부분집합을 새로 채우기 시작하므로, 값을 target으로 나눈 나머지로 저장한다.

즉 dp[mask] = usedSum % target이라고 생각할 수 있다. 도달할 수 없는 상태는 -1로 둔다.

1. 접근 : 문제를 단순화 하기

문제는 정수 배열 nums를 k개의 비어 있지 않은 부분집합으로 나누되, 모든 부분집합의 합이 같게 만들 수 있는지 묻는다. k는 nums.length 이하이고, nums.length는 최대 16이다.

모든 부분집합을 직접 만들어보면 경우의 수가 매우 커진다. 하지만 원소 수가 16개로 작다는 점을 이용하면, "어떤 원소를 이미 골랐는가"를 상태로 삼는 DP가 가능하다.

먼저 전체 합을 sum이라고 하자.

const total = nums.reduce((acc, cur) => acc + cur, 0);

if (total % k !== 0) {
  return false;
}

const target = total / k;

이제 문제는 nums의 모든 원소를 사용하면서, 매번 target을 넘지 않게 숫자를 추가할 수 있는지 확인하는 문제로 바뀐다.

2. 핵심 아이디어

dp[mask]가 -1이 아니라면, 그 상태는 실제로 만들 수 있는 상태다. 이 상태에서 아직 사용하지 않은 원소 nums[i]를 하나 더 넣어본다.

현재 부분집합의 합이 dp[mask]이고, 다음 숫자가 nums[i]라면 다음 합은 다음과 같다.

next = dp[mask] + nums[i]

이때 next > target이면 현재 부분집합이 목표 합을 초과하므로 그 선택은 버린다. 반대로 next <= target이면 다음 마스크로 이동할 수 있다.

다음 상태의 값은 next % target으로 저장한다.

• next < target이면 아직 같은 부분집합을 채우는 중이다.
• next === target이면 부분집합 하나가 완성되었으므로 나머지가 0이 되고, 다음 부분집합을 새로 시작한다.

결국 모든 원소를 사용한 마스크 fullMask에 도달할 수 있고, 그때 나머지가 0이면 모든 부분집합이 정확히 target으로 끝난 것이다.

3. 단계별로 구현하기

3-1. 목표 합 계산과 빠른 예외 처리

전체 합이 k로 나누어떨어지지 않으면 같은 합의 부분집합을 만들 수 없다. 또한 가장 큰 숫자가 target보다 크면 그 숫자를 담을 수 있는 부분집합이 없다.

const total = nums.reduce((acc, cur) => acc + cur, 0);

if (total % k !== 0) {
  return false;
}

const target = total / k;

nums.sort((a, b) => b - a);

if (nums[0] > target) {
  return false;
}

내림차순 정렬은 필수는 아니지만, 큰 수부터 확인하면 불가능한 선택을 더 빨리 걸러낼 수 있다.

3-2. 비트마스크 DP 배열 준비하기

원소가 n개라면 가능한 사용 상태는 2^n개다. 아무 원소도 사용하지 않은 상태 0은 만들 수 있으므로 dp[0] = 0으로 시작한다.

const n = nums.length;
const size = 1 << n;
const dp = new Array(size).fill(-1);

dp[0] = 0;

dp[mask] === -1이면 아직 도달할 수 없는 상태이므로 전이를 시도하지 않는다.

3-3. 아직 쓰지 않은 숫자를 하나씩 추가하기

각 마스크에서 아직 사용하지 않은 원소를 골라본다. 목표 합을 넘지 않는 경우에만 다음 마스크를 reachable 상태로 만든다.

for (let mask = 0; mask < size; mask++) {
  if (dp[mask] === -1) {
    continue;
  }

  for (let i = 0; i < n; i++) {
    if ((mask & (1 << i)) !== 0) {
      continue;
    }

    const next = dp[mask] + nums[i];

    if (next > target) {
      continue;
    }

    const nextMask = mask | (1 << i);
    dp[nextMask] = next % target;
  }
}

다만 같은 nextMask가 이미 도달 가능한 상태라면 다시 계산할 필요가 없다. 그래서 최종 구현에서는 dp[nextMask] !== -1인 경우를 건너뛴다.

4. 최종 구현

/**
 * @param {number[]} nums
 * @param {number} k
 * @return {boolean}
 */
var canPartitionKSubsets = function(nums, k) {
  const total = nums.reduce((acc, cur) => acc + cur, 0);

  if (total % k !== 0) {
    return false;
  }

  const target = total / k;

  nums.sort((a, b) => b - a);

  if (nums[0] > target) {
    return false;
  }

  const n = nums.length;
  const size = 1 << n;
  const dp = new Array(size).fill(-1);

  dp[0] = 0;

  for (let mask = 0; mask < size; mask++) {
    if (dp[mask] === -1) {
      continue;
    }

    for (let i = 0; i < n; i++) {
      if ((mask & (1 << i)) !== 0) {
        continue;
      }

      const next = dp[mask] + nums[i];

      if (next > target) {
        continue;
      }

      const nextMask = mask | (1 << i);

      if (dp[nextMask] !== -1) {
        continue;
      }

      dp[nextMask] = next % target;
    }
  }

  return dp[size - 1] === 0;
};

5. 시간복잡도

n을 nums.length라고 하자.

• 시간복잡도: O(n * 2^n)
• 공간복잡도: O(2^n)

가능한 마스크가 2^n개이고, 각 마스크마다 최대 n개의 원소를 추가해보기 때문이다. 이 문제에서는 n <= 16이므로 이 방식이 충분히 가능하다.

참고자료

• LeetCode: https://leetcode.com/problems/partition-to-k-equal-sum-subsets/