카운팅 테크닉
개요
AtCoder에서 문제를 풀다 보면 백준 온라인 저지나 한국의 대학생 대회에서와 다르게 경우의 수, 기댓값, 확률을 묻는 조합론 문제들이 더 자주 등장한다는 사실을 관찰할 수 있습니다.
한 대회에 출제된 문제 절반 이상에서 $\pmod{998244353}$이 등장해서 너무 과하다고 느껴질 때도 있을 정도입니다.
이 글에서는 제가 문제를 풀면서 자주 보았던 유형들을 몇 가지 선정해서 소개해 보려고 합니다.
각 유형마다 이를 해결하는 테크닉을 소개(널리 통용되는 이름이 없다면 적당히 이름을 붙였습니다)하고, 연습 문제로 AtCoder Regular Contest에서 등장했던 문제를 2개씩 선정했습니다. 글이 너무 길어지는 관계로 모든 문제의 풀이를 적지는 않았습니다.
Sum of Product
각각의 경우의 수마다 $X = (X_1, X_2, \cdots, X_n)$가 결정되고 모든 가능한 경우에 대해 $X_1 \times X_2 \times \cdots \times X_n$들의 합을 구하는 문제는, 곱의 합을 구한다고 생각하는 대신에, 각각의 경우마다 추가로 $X_1$개 중에 하나를 고르고, $X_2$개 중에 하나를 고르고, …, $X_n$개 중에 하나를 고르는 경우의 수를 구하는 문제로 바꿔서 생각할 수 있습니다.
연습 문제 1: Dwango Programming Contest 6th C. Cookie Distribution
문제
두 양의 정수 $N,K$와 수열 $a = (a_1, a_2, \cdots, a_K)$가 주어집니다.
$N$명의 아이가 있고 날짜가 $K$일 있습니다. $i$일차에는 $N$명 중에 $a_i$명을 uniformly random하게 고른 뒤에 선택된 아이들에게 쿠키를 하나씩 줄 것입니다.
$i$번째 아이가 $K$일 동안 받은 쿠키의 개수를 $c_i$라 하면, $c_1 \times c_2 \times \cdots \times c_N$의 기댓값에 $\binom{N}{a_1} \times \binom{N}{a_2} \times \cdots \times \binom{N}{a_K}$를 곱한 값을 $10^9 + 7$로 나눈 나머지를 구해야 합니다.
$(1 \leq N \leq 1000; 1 \leq K \leq 20; 1 \leq a_i \leq N)$
풀이
$K$행 $N$열의 격자를 생각합시다. 각 $i$번째 행마다 $a_i$개의 칸을 골라 쿠키를 놓고, 각 열마다 쿠키가 놓인 칸 중 하나를 골라 음료수를 놓는다고 합시다. 이렇게 만들 수 있는 쿠키와 음료수의 배치의 경우의 수는 원래 문제의 정답과 같습니다.
이제 다음과 같은 DP를 설계할 수 있습니다.
$D[k][i]$ : $k$행의 쿠키 배치를 결정할 차례이고 아직까지 음료수를 놓지 않은 열이 $i$개일 때 올바르게 쿠키와 음료수를 배치하는 경우의 수
음료수를 먼저 배치한 다음에 음료수가 놓인 칸에는 반드시 쿠키도 놓여 있어야 한다고 생각하면 올바른 상태 전이를 유도할 수 있습니다.
이번 행에 $j$개의 음료수를 배치한다면, 음료수를 배치하는 경우의 수는 $\binom{i}{j}$이고 쿠키를 배치하는 경우의 수는 $\binom{N-j}{a_k-j}$가 됩니다.
$D[k][i] = \sum_{0 \leq j \leq i} D[k-1][i-j] \cdot \binom{i}{j} \cdot \binom{N-j}{a_k-j}$와 같은 점화식을 세울 수 있습니다.
정답은 $D[K][N]$이 되고, 시간복잡도는 $O(KN^2)$가 됩니다.
코드
코드는 아래와 같습니다.
#include<bits/stdc++.h>
using namespace std;
#define ll long long
#define MOD 1000000007
int N, K;
int A[21];
ll C[1001][1001], D[21][1001];
int main()
{
ios::sync_with_stdio(0); cin.tie(0);
cin.exceptions(ios::badbit | ios::failbit);
cin >> N >> K;
for(int i=1; i<=K; i++)
cin >> A[i];
for(int n=0; n<=N; n++)
{
C[n][0] = 1;
for(int r=1; r<=n; r++)
C[n][r] = (C[n-1][r] + C[n-1][r-1]) % MOD;
}
D[0][0] = 1;
for(int k=1; k<=K; k++)
for(int i=0; i<=N; i++)
for(int j=0; j<=min(i,A[k]); j++)
(D[k][i] += D[k-1][i-j] * C[i][j] % MOD * C[N-j][A[k]-j]) %= MOD;
cout << D[K][N] << "\n";
return 0;
}
연습 문제 2: AtCoder Regular Contest 157 C. YY Square
mod 2 Lucas
Lucas’s theorem
Lucas’s theorem은 음이 아닌 두 정수 $n,r$과 소수 $p$가 주어질 때, $n$이 매우 크더라도 $\binom{n}{r} \pmod p$을 빠르게 구하는 방법을 알려줍니다.
구체적으로 $n$과 $r$을 $p$진법으로 전개해서,
$n = n_k p^k + n_{k-1} p^{k-1} + \cdots + n_1 p + n_0$이라 하고,
$r = r_k p^k + r_{k-1} p^{k-1} + \cdots + r_1 p + r_0$이라 하면,
$\binom{n}{r} \equiv \prod_{i=0}^{k} \binom{n_i}{r_i} \pmod{p}$가 성립합니다.
Lucas’s theorem, $p = 2$
위의 정리에서 $p = 2$일 경우 비트 연산의 관점에서 생각할 수 있습니다.
$\binom{n_i}{r_i}$는 $n_i = 0, r_i = 1$일 때에만 $0$이기 때문에, $0$이 곱해지지 않으려면 $n$과 $r$을 이진수로 표기했을 때 $r$에서 켜진 비트는 $n$에서도 켜져 있어야 합니다.
따라서 $\land$를 bitwise and 연산이라고 할 때 $\binom{n}{r} \equiv 1 \pmod{2}$과 $n \land r = r$은 동치가 됩니다.
연습 문제 1: AtCoder Regular Contest 137 D. Prefix XORs
문제
길이가 $N$인 수열 $A = (A_1, A_2, \cdots, A_N)$과 정수 $M$이 주어집니다.
각각의 $k = 1,2,\cdots,M$에 대해, 아래 연산을 $k$번 진행했을 때의 $A_N$값을 구해야 합니다.
- 연산: 모든 $1 \leq i \leq N$에 대해 동시에, $A_i$값을 $A_1 \oplus A_2 \oplus \cdots \oplus A_i$로 대체합니다. ($\oplus$는 bitwise xor 연산)
$(1 \leq N,M \leq 10^6; 0 \leq A_i < 2^{30})$
풀이
각 $A_i$가 답에 몇 번이나 XOR되는지, $\textrm{res}[k]$에 대한 $A_i$의 기여도를 생각해 봅시다.
격자에서 생각하면, $A_i$는 $1$행 $i$열에서 오른쪽 이동과 아래 이동만을 사용해서 $k$행 $N$열로 가는 경우의 수, 즉 $\binom{k-1 + N-i}{k-1}$번만큼 기여하게 됩니다.
XOR은 두 번 할 때마다 상쇄되므로 $\binom{k-1 + N-i}{k-1} \equiv 1 \pmod 2$인 경우에만 $\textrm{res}[k]$에 $A_i$가 기여됩니다.
Lucas’s theorem을 적용해서 다시 풀어 적으면 $k-1 + N-i \land k-1 = k-1$일 경우에만 $\textrm{res}[k]$에 $A_i$가 기여됩니다.
받아올림이 일어나지 않아야 하므로, 다시 풀어 적으면 $k-1$과 $N-i$의 이진수 표현이 공유하는 비트가 없는 경우에만 $\textrm{res}[k]$에 $A_i$가 기여됩니다.
마지막으로 다시 풀어 적으면, $k-1$의 모든 비트를 반전한 값의 켜진 비트들의 집합이 $N-i$에서 켜진 비트들의 집합을 subset으로 가지는 경우에만 $\textrm{res}[k]$에 $A_i$가 기여됩니다.
이제 SOS DP로 $O(\max(N,M) \log \max(N,M))$에 해결할 수 있습니다.
SOS DP에 대해서 모르신다면 삼성소프트웨어멤버십 블로그에 있는 다음 글에서 공부하실 수 있습니다.
코드
코드는 아래와 같습니다.
#include<bits/stdc++.h>
using namespace std;
#define INF 1234567890
#define ll long long
int N, M;
int S[1048576];
int main()
{
ios::sync_with_stdio(0); cin.tie(0);
cin.exceptions(ios::badbit | ios::failbit);
cin >> N >> M;
for(int i=1; i<=N; i++)
cin >> S[N-i];
for(int i=0; i<20; i++)
for(int j=0; j<(1<<20); j++)
if (j & 1<<i)
S[j] ^= S[j^1<<i]; // SOS DP
for(int k=1; k<=M; k++)
cout << S[(1<<20)-1 ^ k-1] << " ";
cout << "\n";
return 0;
}
연습 문제 2: AtCoder Regular Contest 156 D. Xor Sum 5
Do Nothing
초기에 원소가 $n$개 있고, 매번 현재 남아 있는 원소들 중에 하나를 uniformly random하게 골라서 삭제하는 연산을 더 이상 진행할 수 없을 때까지 반복할 때, 종료까지 걸리는 연산 횟수의 기댓값을 구하고 싶다고 합시다. 이 때 어떤 연산의 결과로 다른 원소들이 같이 삭제될 수도 있다고 한다고 합시다.
이러한 문제의 경우, 아래처럼 바꿔서 생각해도 연산 횟수의 기댓값은 동일합니다.
- 길이가 $n$인 permutation $n!$개 중에 하나를 uniformly random하게 뽑고, 뽑은 permutation을 $P = (P_1, P_2, \cdots, P_n)$이라 하면 $P$의 원소들을 앞에서부터 차례대로 보면서 $i$번째 차례일 때 $P_i$가 현재 남아 있는 원소라면 골라서 연산을 진행하고, 이미 삭제된 원소라면 아무 행동도 하지 않습니다. (do nothing)
이처럼 다른 연산에 의해 삭제된 원소도 고를 수 있도록 쓸모없는 행동을 추가해도, 연산이 진행될 때에는 매번 현재 남아 있던 원소들 중에 하나가 uniformly random하게 선택되기 때문에 기존과 동일한 조건이 됩니다.
이렇게 문제를 변형한다면 permutation의 관점에서 생각할 수 있고, 이를 이용해서 문제를 더 쉽게 해결할 수도 있습니다.
연습 문제 1: AtCoder Regular Contest 114 E. Paper Cutting 2
문제
$H \times W$ 격자가 주어집니다. 이 격자의 두 칸 $(h_1,w_1)$과 $(h_2,w_2)$은 검은색으로 칠해져 있고, 다른 모든 칸은 흰색으로 칠해져 있습니다.
다음과 같은 과정을 반복한다고 생각해 봅시다.
- 연산: 현재 격자가 $h \times w$ 크기이면 축에 평행한 수평선 $h-1$개와 수직선 $w-1$개가 있을 것입니다. 이 $h+w-2$개의 선들 중 하나를 uniformly random하게 골라서 선을 따라서 격자를 잘라 두 조각으로 만듭니다. 이후,
- 만약 두 검은 칸이 서로 같은 조각에 위치하게 된다면, 검은 칸이 없는 조각을 갖다 버리고 남은 하나의 조각을 새로운 격자로 설정해서 계속 진행합니다.
- 그렇지 않고 두 검은 칸이 서로 다른 조각에 위치하게 된다면, 종료합니다.
위 과정을 종료할 때까지 진행되는 연산 횟수의 기댓값 $\pmod{998244353}$을 구해야 합니다.
$(1 \leq H,W \leq 10^5; HW \geq 2; 1 \leq h_1,h_2 \leq H; 1 \leq w_1,w_2 \leq W; (h_1,w_1) \neq (h_2,w_2))$
풀이
문제를 다른 방식으로 서술하겠습니다.
수평선 $H-1$개와 수직선 $W-1$개의 순서를 섞어서 만들어지는, 길이가 $H+W-2$인 permutation $(H+W-2)!$개 중에 하나를 uniformly random하게 뽑습니다. 뽑은 permutation을 $P = (P_1, P_2, \cdots, P_{H+W-2})$라 하면 $P$의 원소들을 앞에서부터 차례대로 보면서 $i$번째 차례일 때 선 $P_i$가 현재 격자에 존재하는 선이면 연산을 진행하고, 격자에 존재하지 않는 선이면 아무 행동도 하지 않습니다.
이렇게 문제를 변형해도, 종료하기까지의 연산 횟수의 기댓값은 원래 문제의 정답과 같습니다.
이제 각 선이 답에 기여하는 정도를 구해서 모두 더하면 정답을 얻을 수 있습니다.
각 선 $L$에 대해, $L$이 격자에 남아 있기 위해 이전에 고르면 안 되는 선의 개수 $x$는 간단하게 구할 수 있고, $L$이 다른 $x$개의 선보다 permutation 상에서 앞에 위치해야 하므로 $\frac{1}{x+1}$의 확률로 답에 $1$을 기여하게 됩니다.
따라서 각 선마다 $x$값을 구해 주면 $\frac{1}{x+1}$들의 합이 정답이 됩니다.
코드
코드는 아래와 같습니다.
#include<bits/stdc++.h>
using namespace std;
#define ll long long
#define MOD 998244353
ll pw(ll a, ll n)
{
ll ret = 1;
while(n)
{
if (n&1) ret=ret*a%MOD;
a=a*a%MOD;
n>>=1;
}
return ret;
}
ll inv(ll a)
{
return pw(a, MOD-2);
}
int H, W, h1, w1, h2, w2;
int main()
{
ios::sync_with_stdio(0); cin.tie(0);
cin.exceptions(ios::badbit | ios::failbit);
cin >> H >> W >> h1 >> w1 >> h2 >> w2;
if (h1 > h2) swap(h1, h2);
if (w1 > w2) swap(w1, w2);
int ban = h2-h1 + w2-w1;
ll res = 1;
for(int i=0; i<h1-1; i++) res += inv(1+i+ban); // up
for(int i=0; i<w1-1; i++) res += inv(1+i+ban); // left
for(int i=0; i<H-h2; i++) res += inv(1+i+ban); // down
for(int i=0; i<W-w2; i++) res += inv(1+i+ban); // right
res %= MOD;
cout << res << "\n";
return 0;
}
연습 문제 2: AtCoder Regular Contest 165 E. Random Isolation
Inclusion-Exclusion Principle
Inclusion-Exclusion Principle(PIE, 포함 배제의 원리)는 집합들의 교집합의 크기를 이용해서 합집합의 크기를 구하는 방법을 알려줍니다. 집합을 이용해 서술하는 일반적인 방법은 검색으로 쉽게 찾아볼 수 있으니 이 글에서는 제가 생각한 방식대로 적어 보겠습니다.
Inclusion-Exclusion Principle은 주어진 모든 조건을 위반하지 않는 경우의 수를 구하고 싶은 경우에 주로 사용됩니다.
예를 들어, 장애물들의 집합을 $S$라 하면, (어떤 장애물도 지나가지 않고 시작점에서 끝점까지 가는 경우의 수) $= \sum_{i=0}^{\lvert S \rvert} (-1)^{i} \sum_{X \in \mathcal{P}(S), \lvert X \rvert = i}$ (집합 $X$ 내의 모든 장애물을 지나가며 시작점에서 끝점까지 가는 경우의 수) 가 성립합니다.
0개 이상의 장애물을 지나가는 경우의 수를 더해 주고, 1개 이상의 장애물을 지나가는 경우의 수를 빼 주고, 2개 이상의 장애물을 지나가는 경우의 수를 더해 주고, … 이렇게 $\lvert S \rvert$번째까지 번갈아 가면서 더하고 빼 주면 정확히 0개의 장애물을 지나가는 경우의 수를 얻을 수 있습니다.
증명: 정확히 $k$개의 장애물을 지나가는 경우의 수는 위 수식의 $i$번째 스텝 동안 $\binom{k}{i}$번씩 고려되기 때문에 모두 합치면 $\binom{k}{0} - \binom{k}{1} + \binom{k}{2} + \cdots \pm \binom{k}{k} = (1-1)^k$번 세어지게 되고, $k = 0$일 때 $1$번 세어지고 $k > 0$일 때 $0$번 세어지게 되니 정확히 $0$개의 장애물을 지나는 경우의 수를 구하게 됩니다.
위 수식 자체는 지수적인 계산이 필요하지만, 실제 문제 풀이에서는 DP 등의 방법으로 다항 시간 안에 계산할 수 있는 경우가 많습니다.
연습 문제 1: AtCoder Regular Contest 160 D. Mahjong
문제
세 정수 $N,K,M$이 주어지면, 아래의 연산들을 원하는 만큼 자유롭게 사용해서 모든 원소를 $0$으로 만들 수 있는, 음이 아닌 정수들로 이루어진 길이가 $N$이고 합이 $M$인 수열 $A = (A_1, A_2, \cdots, A_N)$의 개수를 $998244353$으로 나눈 나머지를 구해야 합니다.
-
연산 1. $1 \leq i \leq N$인 $i$를 골라서 $A_i$에 $K$를 뺀다.
-
연산 2. $1 \leq i \leq N-K+1$인 $i$를 골라서 $A_i, A_{i+1}, \cdots, A_{i+K-1}$에 $1$을 뺀다.
$(1 \leq K \leq N \leq 2000; 1 \leq M \leq 10^{18})$
풀이
우선, 연산 1과 연산 2는 모두 sum을 $K$씩 감소시키므로 $M$이 $K$의 배수가 아니라면 조건을 만족하는 $A$는 존재하지 않습니다. 이제 $M := M/K$를 한 뒤에 연산을 정확히 $M$번 진행해야 한다고 하겠습니다.
거꾸로 생각해서 $A = (0, 0, \cdots, 0)$에서 출발해서 더하기 연산 $M$번으로 만들 수 있는 $A$의 개수를 구해도 됩니다.
연산을 하는 순서는 상관이 없으므로, 연산 2를 먼저 사용하고 나서야 연산 1을 사용할 수 있도록 강제하겠습니다.
하나의 $i$에다 연산 2를 $K$번 사용하는 것은 $i, i+1, \cdots, i+K-1$에 각각 연산 1을 한 번씩 사용하는 것으로 대체할 수 있으므로, 하나의 $i$에는 연산 2를 $K$번 미만으로만 사용할 수 있다고 강제해도 됩니다.
이제 $A_i \bmod K$값을 바꿀 수 있는 건 연산 2뿐이므로, 어떤 두 방법에서 연산 2를 모두 사용한 이후의 수열 $A$가 서로 다르다면, 이후에 연산 1을 어떻게 사용하더라도 $A$가 다시 같아지는 방법은 존재하지 않습니다.
$i$에 연산 1을 사용한 횟수를 $X_i$라 하고 연산 2를 사용한 횟수를 $Y_i$라 하겠습니다.
-
$1 \leq i \leq N$에서 $X_i \geq 0$
-
$1 \leq i \leq N-K+1$에서 $0 \leq Y_i < K$
-
$\sum X_i + \sum Y_i = M$
이제 위의 세 조건을 모두 만족하는 $(X,Y)$의 개수를 세는 문제가 되었습니다.
$Y_i$ 값은 $K$ 미만이어야 한다. 라는 조건이 $N-K+1$개 있는 것으로 생각할 수 있으니 Inclusion-Exclusion Principle을 적용할 수 있습니다.
조건들의 집합을 $S$라 하면, (어떤 조건도 위반하지 않는 $(X,Y)$의 개수) $= \sum_{i=0}^{N-K+1} (-1)^{i} \sum_{s \in \mathcal{P}(S), \lvert s \rvert = i}$ (집합 $s$ 내의 모든 조건을 위반하는 $(X,Y)$의 개수) 가 됩니다.
위반해야 하는 조건들에 대응되는 $Y_i$값들에 먼저 $K$씩을 더해 주고, 그 뒤에 남은 $M - \lvert s \rvert \cdot K$을 분배한다고 합시다.
$M - \lvert s \rvert \cdot K$개의 공을 $2N-K+1$개의 상자에 중복을 허용해서 분배하는 방법의 수는 중복조합을 사용하면 $\binom{2N-K + M - \lvert s \rvert \cdot K}{2N-K}$가 됩니다. 또한 집합 $s$의 크기가 같으면 내부 구성 원소는 중요하지 않습니다.
이제 (어떤 조건도 위반하지 않는 $(X,Y)$의 개수) $= \sum_{i=0}^{N-K+1} (-1)^{i} \binom{N-K+1}{i} \binom{2N-K + M-iK}{2N-K}$ 가 됩니다.
$M \leq 10^{18}$로 매우 크지만 $2N-K$는 작기 때문에 매번 $O(N)$ 시간에 이항 계수를 계산할 수 있고, 최종 시간복잡도는 $O(N^2)$가 됩니다.
코드
$M$이 매우 크기 때문에 조심해서 나머지 연산을 해야 합니다.
코드는 아래와 같습니다.
#include<bits/stdc++.h>
using namespace std;
#define INF 1234567890
#define ll long long
#define MOD 998244353
ll pw(ll a, ll n)
{
ll ret = 1;
while(n)
{
if (n&1) ret=ret*a%MOD;
a=a*a%MOD;
n>>=1;
}
return ret;
}
ll memo[4001];
ll inv(ll a)
{
if (memo[a]) return memo[a];
return memo[a] = pw(a, MOD-2);
}
ll C(ll n, ll r)
{
if (n < 0 || r < 0 || r > n) return 0LL;
ll ret = 1;
for(int i=0; i<r; i++)
{
(ret *= (n-i) % MOD) %= MOD;
(ret *= inv(i+1)) %= MOD;
}
return ret;
}
int N, K;
ll M;
int main()
{
ios::sync_with_stdio(0); cin.tie(0);
cin.exceptions(ios::badbit | ios::failbit);
cin >> N >> M >> K;
if (M % K) { cout << "0\n"; return 0; }
M /= K;
ll res = 0;
for(int i=0; i<=N-K+1; i++)
(res += (i % 2 ? -1 : 1) * C(N-K+1, i) * C(2*N-K + M-i*K, 2*N-K)) %= MOD;
res += MOD; res %= MOD;
cout << res << "\n";
return 0;
}