백준 23824번 서버 증축

문제

백준 23824번 서버 증축

풀이

정수 $k, a, n$이 주어진다. $(1 \le k \le 40;$ $1 \le a \le {10}^{6};$ $1 \le n \le {10}^{15})$

중복집합 $A = \{\underbrace{2^0, 2^0, \cdots, 2^0}_{a\text{ times}}, \underbrace{2^1, 2^1, \cdots, 2^1}_{a\text{ times}}, \cdots, \underbrace{2^{k-1}, 2^{k-1}, \cdots, 2^{k-1}}_{a\text{ times}}\}$의 부분집합을 $B$라고 할 때 $\sum_{x \in B}{x}=n$인 $B$의 개수를 구해야 한다.

다음과 같은 생성함수 $F(x)$를 생각하자.

$$
\begin{align}
F(x)
&= \underbrace{\left(1+x^{2^0}\right)\left(1+x^{2^0}\right)\cdots\left(1+x^{2^0}\right)}_{a\text{ times}} \underbrace{\left(1+x^{2^1}\right)\left(1+x^{2^1}\right)\cdots\left(1+x^{2^1}\right)}_{a\text{ times}} \cdots \underbrace{\left(1+x^{2^{k-1}}\right)\left(1+x^{2^{k-1}}\right)\cdots\left(1+x^{2^{k-1}}\right)}_{a\text{ times}} \\
&= \prod_{i=0}^{k-1}{\left(1 + x^{2^i}\right)^{a}}
\end{align}
$$

이때 $[x^n]F(x)$가 답이다. $[x^n]F(x)$는 $F(x)$에서 $x^n$의 계수이다. $[x^n]F(x)$가 합이 $n$인 부분집합의 개수라는 것에 대해서는 3B1B의 영상 Olympiad level counting: How many subsets of {1,…,2000} have a sum divisible by 5?을 보면 좋다.

$F(x)$를 다음과 같이 정리할 수 있다.

$$
\begin{align}
F(x) & = \prod_{i=0}^{k-1}{\left(1 + x^{2^i}\right)^{a}} \\
& = \left( \prod_{i=0}^{k-1}{\left(1 + x^{2^i}\right)}\right)^{a} \\
& = \left(x^0 + x^1 + \cdots + x^{2^{k}-1}\right)^{a} \\
& = \left(\sum_{i = 0}^{2^{k}-1}{x^i}\right)^{a}
\end{align}
$$

이때 $\prod_{i=0}^{k-1}{\left(1 + x^{2^i}\right)} = \sum_{i = 0}^{2^{k}-1}{x^i}$가 성립하는데 이는 집합 $\{2^0, 2^1, \cdots, 2^{k-1}\}$에서 $0 \le i \le 2^k-1$인 정수 $i$에 대해 합이 $i$가 되도록 원소를 뽑는 경우가 유일하게 존재하기 때문이다. 이는 $i$의 이진법 표현을 생각해보면 쉽게 알 수 있다.

이제 $[x^n]\left(\sum_{i = 0}^{2^{k}-1}{x^i}\right)^{a}$를 구해야 한다. 입력 조건 때문에 뤼카의 정리를 사용해야 한다는 것을 제외하면 이는 16725번 다항 계수 문제와 동일하다.

결론적으로 $[x^n]F(x) = \sum_{i=0}^{\min\left(a, \lfloor\frac{n}{2^k}\rfloor\right)}{\left(-1\right)^{i}\binom{a}{i}\binom{a+n-2^{k}i-1}{a-1}}$이다. 이는 포함배제의 원리를 통해서 보이거나 생성함수를 이용해서 보일 수 있다.

포함배제의 원리. $[x^n]\left(\sum_{i = 0}^{2^{k}-1}{x^i}\right)^{a}$는 $x^n = x^{\alpha_1}x^{\alpha_2} \cdots x^{\alpha_a}$인 $\left(\alpha_1, \alpha_2, \cdots, \alpha_a\right)$의 개수이다. $(0 \le \alpha \le 2^k - 1;$ $\alpha_i\text{는 정수})$ $x^{\alpha_1}x^{\alpha_2} \cdots x^{\alpha_a} = x^{\alpha_1 + \alpha_2 + \cdots + \alpha_a}$이므로 $\alpha_1 + \alpha_2 + \cdots + \alpha_a = n$를 만족하는 $\left(\alpha_1, \alpha_2, \cdots, \alpha_a\right)$의 개수를 구하면 된다.

$0 \le \alpha_i$이면 $\alpha_1 + \alpha_2 + \cdots + \alpha_a = n$의 경우의 수는 $\left(\!\!{a \choose n}\!\!\right)$이다. $\left(\!\!{n \choose k}\!\!\right)$는 중복 조합의 수이다. 이를 이용해서 $0 \le \alpha_i \le 2^k-1$일 때의 경우의 수를 포함배제의 원리로 다음과 같이 나타낼 수 있다.

$$
\sum_{i=0}^{a}{\left(-1\right)^{i} \binom{a}{i} \left(\!\!{{a} \choose {n - 2^{k}i}}\!\!\right)}
= \sum_{i=0}^{\min\left(a, \lfloor\frac{n}{2^k}\rfloor\right)}{\left(-1\right)^{i} \binom{a}{i} \binom{a + n - 2^{k}i - 1}{a - 1}}
$$

$0 \le i \le \min\left(a, \lfloor\frac{n}{2^k}\rfloor\right)$는 단순히 이항 계수가 $0$이 아닐 조건이다. 그리고 $\left(\!\!{n \choose k}\!\!\right) = \binom{n + k - 1}{k}$이다.

따라서 $[x^n]F(x) = \sum_{i=0}^{\min\left(a, \lfloor\frac{n}{2^k}\rfloor\right)}{\left(-1\right)^{i} \binom{a}{i} \binom{a + n - 2^{k}i - 1}{a - 1}}$이다.

생성함수. $F(x) = \left(\sum_{i = 0}^{2^{k}-1}{x^i}\right)^{a}$의 형태를 변형하여 $[x^n]F(x)$를 구한다. 이때 아래의 사실들을 활용한다.

• $\sum_{i = 0}^{n-1}{x^i} = \frac{1-x^{n}}{1-x}$
• $\sum_{i = 0}^{n}{\left(-1\right)^{i}\binom{n}{i} x^i} = \left(1-x\right)^n$
• $\sum_{i = 0}^{\infty}{\left(\!\!{n \choose i}\!\!\right)x^i} = \frac{1}{\left(1-x\right)^n}$

$$
\begin{align}
F(x) & = \left(\sum_{i = 0}^{2^{k}-1}{x^i}\right)^{a} \\
& = \left(\frac{1-x^{2^{k}}}{1-x}\right)^a \\
& = \left( 1-x^{2^{k}} \right)^{a} \frac{1}{\left( 1-x \right)^{a}} \\
& = \sum_{i = 0}^{a}{\left(-1\right)^{i}\binom{a}{i}x^{2^{k}i}} \sum_{i = 0}^{\infty}{\left(\!\!{a \choose i}\!\!\right)x^i}
\end{align}
$$

$\sum_{i = 0}^{a}{\left(-1\right)^{i}\binom{a}{i}x^{2^{k}i}} \sum_{i = 0}^{\infty}{\left(\!\!{a \choose i}\!\!\right)x^i}$를 전개한 식을 생각하면 $[x^n]F(x) = \sum_{i=0}^{\min\left(a, \lfloor\frac{n}{2^k}\rfloor\right)}{\left(-1\right)^{i} \binom{a}{i} \left(\!\!{{a} \choose {n - 2^{k}i}}\!\!\right)}$임을 알 수 있다.

따라서 $[x^n]F(x) = \sum_{i=0}^{\min\left(a, \lfloor\frac{n}{2^k}\rfloor\right)}{\left(-1\right)^{i} \binom{a}{i} \binom{a + n - 2^{k}i - 1}{a - 1}}$이다.

이제 $[x^n]F(x) \pmod{1048573}$을 구해야 한다. $[x^n]F(x) = \sum_{i=0}^{\min\left(a, \lfloor\frac{n}{2^k}\rfloor\right)}{\left(-1\right)^{i} \binom{a}{i} \binom{a + n - 2^{k}i - 1}{a - 1}}$이고 $1 \le a \le {10}^{6}$이므로 $0$부터 $\min\left(a, \lfloor\frac{n}{2^k}\rfloor\right)$까지의 $i$에 대해 $\left(-1\right)^{i} \binom{a}{i} \binom{a + n - 2^{k}i - 1}{a - 1} \pmod{1048573}$을 직접 계산하는 것으로 충분하다.

$1 \le n \le {10}^{15}$이고 $1048573$은 소수이므로 이항 계수를 계산할 때는 뤼카의 정리를 이용한다.

코드

#include <atcoder/modint>
#include <bits/stdc++.h>
using namespace std;
constexpr auto mod = (1 << 20) - 3;
using mint = atcoder::static_modint<mod>;

int main() {
    cin.tie(0)->sync_with_stdio(0);
    long long k, a, n;
    cin >> k >> a >> n;
    array<mint, mod> f;
    f[0] = 1;
    for (auto i = 1; i < mod; ++i) {
        f[i] = f[i - 1] * i;
    }
    array<mint, mod> finv;
    finv[mod - 1] = f[mod - 1].inv();
    for (auto i = mod - 1; i > 0; --i) {
        finv[i - 1] = finv[i] * i;
    }
    auto C_small = [&](auto m, auto r) {
        if (0 <= r && r <= m) {
            return f[m] * finv[r] * finv[m - r];
        } else {
            return mint();
        }
    };
    auto C = [&](auto m, auto r) {
        assert(0 <= m && 0 <= r);
        mint acc(1);
        while (m > 0) {
            acc *= C_small(m % mod, r % mod);
            m /= mod;
            r /= mod;
        }
        return acc;
    };
    mint ans, g(1);
    for (auto i = 0LL; i <= min(a, n / (1LL << k)); ++i) {
        ans += g * C(a, i) * C(a + n - i * (1LL << k) - 1, a - 1);
        g *= -1;
    }
    cout << ans.val();
}

관련 문제

• 13542번 우표 구매하기
• 16725번 다항 계수

참고 자료

• Olympiad level counting: How many subsets of {1,…,2000} have a sum divisible by 5?
• [Tutorial] Generating Functions in Competitive Programming (Part 1)
• [Tutorial] Generating Functions in Competitive Programming (Part 2)
• Problem Solving과 생성함수
• 다양한 생성함수와 그 응용 (1)
• 다양한 생성함수와 그 응용 (2)