Computer >> 컴퓨터 >  >> 프로그래밍 >> C++

C++에서 배열을 합이 동일한 K개의 부분 집합으로 분할하는 방법

문제 개요

정수 배열 nums와 양의 정수 k가 주어졌을 때, 이 배열을 각 부분 집합의 합이 서로 같도록 k개의 비어 있지 않은 부분 집합으로 나눌 수 있는지 확인하는 것이 목표입니다.

예를 들어 배열이 [4, 3, 2, 3, 5, 2, 1]이고 k = 4라고 가정해 보겠습니다. 이 경우 결과는 True입니다. 배열을 [[5], [1, 4], [2, 3], [2, 3]]처럼 네 개의 그룹으로 나누면 각 그룹의 합이 모두 5로 동일하기 때문입니다.

해결 전략: 비트마스크 동적 계획법(DP)

부분 집합 분할 문제는 가능한 조합의 수가 원소 개수에 따라 지수적으로 증가하기 때문에 단순한 완전 탐색으로는 비효율적입니다. 대신 각 원소의 포함 여부를 비트(bit)로 표현하는 비트마스크 DP 기법을 활용하면 효율적으로 해결할 수 있습니다. n개의 원소에 대해 2^n개의 상태를 테이블로 관리하고, 각 상태마다 누적 합계를 함께 저장하여 유효한 분할만 추적합니다.

알고리즘 단계

  1. 크기가 2^n인 두 개의 테이블 dp(상태 도달 가능 여부)와 total(누적 합계)을 정의합니다.
  2. 배열 nums를 오름차순으로 정렬하고, 모든 원소의 합을 sum에 저장합니다.
  3. sum % k가 0이 아니거나, 가장 큰 원소(nums의 마지막 값)가 sum / k보다 크면 분할이 불가능하므로 즉시 false를 반환합니다.
  4. dp[0] := true로 설정하고, sum := sum / k로 갱신합니다. 이때부터 sum은 각 부분 집합이 가져야 할 목표 합이 됩니다.
  5. i를 0부터 2^n - 1까지 순회하며 다음을 수행합니다.
    • dp[i]가 참이라면 j를 0부터 n - 1까지 반복합니다.
      • temp := i OR 2^j로 j번째 원소를 추가한 새로운 상태를 만듭니다.
      • temp가 i와 다르다면(j번째 원소가 아직 포함되지 않았다면):
        • nums[j] <= sum - (total[i] % sum)을 만족하면 dp[temp] := true로 설정하고, total[temp] := total[i] + nums[j]로 누적 합계를 갱신합니다.
        • 조건을 만족하지 않으면 내부 반복문을 종료(break)합니다.
  6. 최종적으로 dp[(2^n) - 1], 즉 모든 원소를 사용한 상태의 값을 반환합니다.

핵심 아이디어: total[i] % sum은 현재 상태 i에서 마지막으로 완성된 부분 집합 이후 새로 채워야 하는 남은 합을 의미합니다. 다음 원소가 이 남은 공간에 들어갈 수 있을 때만 상태를 확장함으로써, 각 부분 집합의 합이 정확히 sum이 되도록 보장할 수 있습니다. 또한 배열을 미리 오름차순으로 정렬해 두면, 어떤 원소가 남은 공간에 들어가지 못할 때 그보다 큰 뒤의 원소들 역시 들어갈 수 없으므로 반복문을 조기에 종료하는 최적화가 가능합니다.

C++ 구현 예제

아래 구현을 통해 더 자세히 이해해 보겠습니다.

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
    bool canPartitionKSubsets(vector<int>& nums, int k) {
        int n = nums.size();
        vector <bool> dp(1 << n);
        vector <int> total(1 << n);
        sort(nums.begin(), nums.end());
        int sum = 0;
        for(int i = 0; i < nums.size(); i++)sum += nums[i];
        if(sum % k || nums[nums.size() - 1] > sum / k) return false;
        dp[0] = true;
        sum /= k;
        for(int i = 0; i < (1 << n); i++){
            if(dp[i]){
                for(int j = 0; j < n; j++){
                    int temp = i | (1 << j);
                    if(temp != i){
                        if(nums[j] <= sum - (total[i] % sum)){
                           dp[temp] = true;
                           total[temp] = total[i] + nums[j];
                        }
                        else{
                           break;
                        }
                    }
                }
            }
        }
        return dp[(1 << n) - 1];
    }
};
main(){
    Solution ob;
    vector<int> v = {4,3,2,3,5,2,1};
    cout << (ob.canPartitionKSubsets(v, 4));
}

입력

[4,3,2,3,5,2,1]
4

출력

1

출력값 1(true)은 주어진 배열을 합이 5로 동일한 네 개의 부분 집합 [[5], [1, 4], [2, 3], [2, 3]]으로 성공적으로 나눌 수 있음을 의미합니다.