본문으로 건너뛰기

3343 - Count Number of Balanced Permutations

정보

풀이 키워드​

스포주의

DP 수학 조합론


풀이 코드​

정보
  • 메모리: 45300 KB
  • 시간: 96 ms
class Solution {
    final int MOD = 1_000_000_007;
    int[] cnt = new int[10];
    long[] fact = new long[81];
    long[] inv = new long[81];
    long[] invFact = new long[81];
    int n, evenLen, oddLen, halfSum;
    int[][][] memo;

    private long dp(int i, int depth, int curSum) { // i: current digit | depth: length of selected
        if (curSum > halfSum || depth > evenLen) return 0;
        if (i > 9) return (depth == evenLen && curSum == halfSum) ? 1 : 0;
        if (memo[i][depth][curSum] != -1) return memo[i][depth][curSum];

        long res = 0;
        for (int use = 0; use <= cnt[i] && depth+use <= evenLen && curSum+i*use <= halfSum; ++use) {
            long evenPerm = fact[depth+use] * invFact[depth] % MOD * invFact[use] % MOD;
            long oddInv = invFact[cnt[i]-use]; // 안쓴만큼 홀수 인덱스로 감 -> 홀수 인덱스 배치 수식의 분모
            res = (res + (evenPerm * oddInv % MOD) * dp(i+1, depth+use, curSum+i*use)) % MOD;
        }

        memo[i][depth][curSum] = (int) res;
        return res;
    }

    public int countBalancedPermutations(String num) {
        n = num.length();
        int totalSum = 0;
        for (char ch : num.toCharArray()) {
            totalSum += ch - '0';
            ++cnt[ch-'0'];
        }
        if ((totalSum & 1) == 1) return 0; // can't satisfy condition

        halfSum = totalSum / 2;
        evenLen = (n+1) / 2;
        oddLen = n - evenLen;
        memo = new int[10][evenLen+1][halfSum+1];
        for (int i = 0; i < 10; ++i)
            for (int j = 0; j <= evenLen; ++j)
                Arrays.fill(memo[i][j], -1);
        
        precalc();
        long res = dp(0, 0, 0);
        // 홀수 인덱스 배치 수식의 분자 곱하기
        res = res * fact[oddLen] % MOD;
        return (int) res;
    }

    public void precalc() {
        fact[0] = inv[0] = invFact[0] = 1;
        for (int i = 1; i < 81; ++i) fact[i] = fact[i-1] * i % MOD;
        inv[1] = 1; // extended euclidean
        for (int i = 2; i < 81; ++i) inv[i] = MOD - (MOD/i) * inv[MOD%i] % MOD;
        for (int i = 1; i < 81; ++i) invFact[i] = invFact[i-1] * inv[i] % MOD;
    }
}

풀이 해설​

주어진 문자열의 숫자를 뒤섞어서 홀수 인덱스 합, 짝수 인덱스 합이 같도록 하는 경우의 수를 세는 문제이다.

경우의 수가 매우 많을 수 있어 모듈러 연산을 수행하며 중복순열 개수를 계산해야 하는데,

모듈러 연산 공식에 나눗셈은 정의되어있지 않으므로 모듈러 역원을 사용해야 한다.


📌 모듈러 역원 구하기​

모듈러 역원은 페르마의 소정리 or 확장 유클리드 호제법 중 선택하여 구현할 수 있고

만약 사전 계산하여 배열로 다 구해둘 목적이면 총 O(n log n)O(n\,log\,n) 복잡도를 가진 페르마 소정리보단

재귀적으로 처리하는 확장 유클리드 호제법이 좀 더 나을 수 있겠다.

하지만 쓰이는 역원이 한정되어있기 때문에 퍼포먼스를 위해선 페르마로 가는게 나아보인다. ㅡ_ㅡ


📌 백트래킹으로는 부족하다​

풀이 코드의 dp 함수가 원래는 dfs였고 prev 인자를 가지고 있었다.

prev는 이전에 선택한 숫자값이고, 이전 선택 값 미만으로는 선택하지 않게끔 로직을 작성했었다.

하지만 num의 길이가 길어지면 다음의 TC에서 TLE가 터진다.

백트래킹 코드
class Solution {
    final int MOD = 1_000_000_007;
    int[] cnt = new int[10];
    int[] selectedCnt = new int[10];
    long[] fact = new long[41]; // index: 0 ~ 80/2
    long[] inv = new long[41];
    long[] invFact = new long[41];
    int n = 0;
    int sum = 0; // should always be even

    public int calc(int len) {
        long res1 = fact[len];
        for (int i = 0; i < 10; ++i) {
            if (selectedCnt[i] == 0) continue;
            res1 = res1 * invFact[selectedCnt[i]] % MOD;
        }

        long res2 = fact[n-len];
        for (int i = 0; i < 10; ++i) {
            if (cnt[i] == 0) continue;
            res2 = res2 * invFact[cnt[i]] % MOD;
        }
        
        return (int) (res1 * res2 % MOD);
    }

    public int dfs(int prev, int curSum, int depth) { // prev: last selected digit
        if (curSum > sum/2) return 0;
        if (depth == n/2+1) {
            return curSum == sum/2 ? calc(depth-1) : 0;
        }

        int res = 0;
        for (int i = prev; i < 10; ++i) {
            if (cnt[i] == 0) continue; // can't select
            --cnt[i];
            ++selectedCnt[i];
            res = (res + dfs(i, curSum+i, depth+1)) % MOD;
            --selectedCnt[i];
            ++cnt[i];
        }

        return res;
    }

    public int countBalancedPermutations(String num) {
        this.n = num.length();
        for (char ch : num.toCharArray()) {
            ++cnt[ch-'0'];
            sum += (ch-'0');
        }
        
        if ((sum & 1) == 1) return 0; // can't satisfy condition

        precalc();
        return dfs(0, 0, 1);
    }

    public void precalc() {
        fact[0] = inv[0] = invFact[0] = 1;
        for (int i = 1; i < 41; ++i) fact[i] = fact[i-1] * i % MOD;
        inv[1] = 1;
        for (int i = 2; i < 41; ++i) inv[i] = MOD - (MOD/i) * inv[MOD%i] % MOD;
        for (int i = 1; i < 41; ++i) invFact[i] = invFact[i-1] * inv[i] % MOD;
    }
}
"9103224235015855909098892653211286258658354677013618352"

메모​

  • 영미권 문제에서 digit이라고 하면 leading zero가 허용되고, number라고 하면 안된다.
  • DP 정의가 까다로운 편이다.
  • 대회 문제이고 일반적인 코테에 나올 리 없는 유형.
    • 이걸 데일리 문제로 내놓는게 레전드...😮