BOJ 11618:: Frightful Formula

https://www.acmicpc.net/problem/11618

얼핏 보면 단순한 동적계획법 문제이지만 무식하게 풀면 시간복잡도가 $O(n^2)$으로 절대 제한 시간 내에 풀 수 없습니다. 따라서 $F[n,n]$을 $l_k$와 $t_k$로 나타내는 식을 구하고, 그걸 최대한 빠르게 계산하는 방향으로 풀어야 합니다.

그러나 이 $F[n,n]$의 식을 구하는 게 매우 어렵습니다. 생각의 전환이 필요한데, 각 $l_k$, $t_k$가 $F[n,n]$에 기여하는 정도를 계산할 겁니다. 여기서 기여하는 정도라는 건, 어떤 칸 $F[x,y]$의 값을 다른 칸 $F[i, j]$를 포함하는 수식으로 나타냈을 때 $F[i,j]$의 계수를 $F[i,j]$가 $F[x,y]$에 기여하는 정도라고 정의합니다. 예를 들어, $F[i,j]$는 $F[i,j+1]$에 $a$만큼 기여하고, $F[i,j]$는 $F[i+1,j]$에 $b$만큼 기여합니다. 마찬가지로 $F[i,j+2]$에는 $a^2$만큼, $F[i+1,j+1]$에는 $2ab$만큼, $F[i+2,j]$에는 $b^2$만큼 기여합니다. 즉, $F[i,j]$의 기여도는 오른쪽 아래로 내려갈수록 점화식에 의해 증폭되고, 증폭 정도는 오른쪽으로 한 칸 갈 때 $a$배, 아래로 한 칸 갈 때 $b$배입니다. 따라서 $F[i,j]$가 $F[x,y]$에 기여하는 정도는

  • $F[i,j]$는 $F[i,j]$에 1만큼 기여
  • 이 기여도가 $(i,j)$에서 $(x,y)$까지 가면서 $a^{x-i}b^{y-j}$배 증폭
  • 그런데 $(i,j)$에서 $(x,y)$까지 갈 수 있는 경로는 $\binom{x+y-i-j}{x-i}$개 있으므로 이를 곱해주어야 함

예를 들어, $F[i,j]$가 $F[i+1,j+1]$에 기여하는 정도를 구한다고 하면 $F[i,j]$는 $F[i,j]$에 1만큼 기여하고, 이 기여도가 오른쪽으로 한 칸, 아래로 한 칸 이동하면서 $ab$배로 증폭됩니다. 그런데 $(i,j)$에서 $(i+1,j+1)$까지 가는 경로는 두 개 있으므로 증폭이 두 번 됩니다. 따라서 최종 기여도는 $2ab$입니다.

따라서 $l_k$는 $F[k,2]$에 $a$만큼 기여하고, $(n,n)$까지 가면서 기여도는 $\binom{2n-k-2}{n-2}a^{n-2}b^{n-k}$배로 증폭되므로 $l_k$가 $F[n,n]$에 기여하는 정도는 $\binom{2n-k-2}{n-2}a^{n-1}b^{n-k}$가 됩니다. 같은 방법으로 $t_k$는 $F[n,n]$에 $\binom{2n-k-2}{n-2}a^{n-k}b^{n-1}$만큼 기여합니다. 이 기여도를 다 더하면

\[ F[n,n] = \sum_{i=2}^n l_i \binom{2n-i-2}{n-2}a^{n-1}b^{n-i} + \sum_{i=2}^n t_i \binom{2n-i-2}{n-2}a^{n-i}b^{n-1} \]

뭔가 이상하죠? $c$가 없네요. $F[i,j]$의 점화식에선 $F[i,j-1]$과 $F[i-1,j]$에 포함된 $l_k$와 $t_k$의 기여도를 증폭하는 것 말고도 새로 $c$ 항이 더해집니다. 이 $c$도 $(n,n)$까지 가면서 똑같이 증폭됩니다. 그러므로,

\[ F[n,n] = \sum_{i=2}^n l_i \binom{2n-i-2}{n-2}a^{n-1}b^{n-i} + \sum_{i=2}^n t_i \binom{2n-i-2}{n-2}a^{n-i}b^{n-1}
+ c \sum_{i=2}^n \sum_{j=2}^n \binom{2n-i-j}{n-i} a^{n-j} b^{n-i} \]

$0!$부터 $(2n)!$까지와 $a$, $b$의 거듭제곱을 미리 계산해두면 각 이항 계수는 바로 구할 수 있으므로 앞 두 항은 $O(n)$에 구할 수 있습니다. 마지막 항은 그냥 구하면 $O(n^2)$이지만 저걸 잘 바꿔서 합성곱 형태로 만들어서 FFT로 $O(n\log n)$에 구할 수 있습니다.

… 이게 대회 공식 풀이고, 사실 훨씬 간단한 방법으로 풀 수 있습니다.

저 마지막 항은 아래와 같이 고칠 수 있습니다.

\[ c \sum_{i=0}^{n-2} \sum_{j=0}^{n-2} \binom{i+j}{i} a^j b^i \]

$f(n) = \sum_{i=0}^n \sum_{j=0}^n \binom{i+j}{i} a^j b^i$으로 정의하면 저 값은 $cf(n-2)$입니다. 그리고

\[ \begin{align}
f(n) &= \sum_{i=0}^n \sum_{j=0}^n \binom{i+j}{i} a^j b^i \\
&= \underbrace{\sum_{i=0}^{n-1} \sum_{j=0}^{n-1} \binom{i+j}{i} a^j b^i}_{f(n-1)}
+ a^n \sum_{i=0}^n \binom{n+i}{n} b^i + b^n \sum_{i=0}^n \binom{n+i}{n} a^i – \binom{2n}{n} a^n b^n
\end{align} \]

또 $g_a(n) = \sum_{i=0}^n \binom{n+i}{n} a^i$로 정의하면

\[ \begin{align}
g_a(n) &= \sum_{i=0}^n \binom{n+i}{n} a^i \\
&= a + \sum_{i=1}^n \binom{n+i}{n} a^i \\
&= a + \sum_{i=1}^n \left[ \binom{n+i-1}{n-1} + \binom{n+i-1}{n} \right] a^i \\
&= a + \sum_{i=1}^{n-1} \binom{n+i-1}{n-1} a^i + \binom{2n-1}{n-1} a^n + \sum_{i=1}^n \binom{n+i-1}{n} a^i \\
&= \sum_{i=0}^{n-1} \binom{n+i-1}{n-1} a^i + \binom{2n-1}{n-1} a^n + \sum_{i=0}^{n-1} \binom{n+i}{n} a^{i+1} \\
&= \underbrace{\sum_{i=0}^{n-1} \binom{n+i-1}{n-1} a^i}_{g_a(n-1)} + \binom{2n-1}{n-1} a^n +
a \underbrace{\sum_{i=0}^n \binom{n+i}{n} a^i}_{g_a(n)} – \binom{2n}{n} a^{n+1}
\end{align} \]

\[ \therefore g_a(n) = \frac{1}{1-a} \left[ g_a(n-1) + \binom{2n-1}{n} a^n – \binom{2n}{n} a^{n+1} \right] \]

단, 여기서 $a=1$인 경우에는 위 식에서

\[ g_a(n-1) = \binom{2n}{n} – \binom{2n-1}{n} = \binom{2n-1}{n-1} \]

\[ \therefore g_a(n) = \binom{2n+1}{n} \]

마찬가지로 $g_b(n) = \sum_{i=0}^n \binom{n+i}{n} b^i$으로 정의하면 $a$ 대신 $b$를 대입한 식이 나옵니다. 정리하자면

\[ \begin{align}
g_a(n) &= \begin{cases}
\frac{1}{1-a} \left[ g_a(n-1) + \binom{2n-1}{n} a^n – \binom{2n}{n} a^{n+1} \right] & \textrm{if } a \neq 1 \\
\binom{2n+1}{n} & \textrm{if } a = 1
\end{cases} \\
g_b(n) &= \begin{cases}
\frac{1}{1-b} \left[ g_b(n-1) + \binom{2n-1}{n} b^n – \binom{2n}{n} b^{n+1} \right] & \textrm{if } b \neq 1 \\
\binom{2n+1}{n} & \textrm{if } b = 1
\end{cases} \\
f(n) &= f(n-1) + a^n g_b(n) + b^n g_a(n) – \binom{2n}{n} a^n b^n
\end{align} \]

$g_a(n)$, $g_b(n)$, $f(n)$은 모두 $O(n)$에 계산 가능하고, 따라서 $cf(n-2)$도 $O(n)$에 계산할 수 있습니다.

#include<bits/stdc++.h>


using namespace std;

long long n, a, b, c;
long long l[200001], t[200001];
long long powa[200001] = {1}, powb[200001] = {1};     // powa[i] = a^i, powb[i] = b^i
long long fac[400001] = {1}, facinv[400001];          // fac[i] = i!, facinv[i] = 1/i!
long long f[200001] = {1}, ga[200001] = {1}, gb[200001] = {1};

const long long MOD = 1000003;

// calculate x^y % MOD
inline long long powmod(long long x, long long y) {
    if (y == 0)
        return 1;
    long long half = powmod(x, y/2);
    if (y % 2 == 1)
        return half * half % MOD * x % MOD;
    return half * half % MOD;
}

// calculate binomial coefficient Comb(x, y)
inline long long binom(long long x, long long y) {
    return fac[x] * facinv[y] % MOD * facinv[x-y] % MOD;
}


int main(void) {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);

    cin >> n >> a >> b >> c;
    for (int i = 1; i <= n; i++)
        cin >> l[i];
    for (int i = 1; i <= n; i++)
        cin >> t[i];

    for (int i = 1; i <= n; i++) {
        powa[i] = powa[i-1] * a % MOD;
        powb[i] = powb[i-1] * b % MOD;
    }
    for (int i = 1; i <= 2*n; i++)
        fac[i] = fac[i-1] * i % MOD;
    facinv[2*n] = powmod(fac[2*n], MOD-2);
    for (int i = 2*n-1; i >= 0; i--)
        facinv[i] = facinv[i+1] * (i+1) % MOD;

    long long inv1a = powmod(1-a+MOD, MOD-2), inv1b = powmod(1-b+MOD, MOD-2);
    for (int i = 1; i <= n-2; i++) {
        ga[i] = (a == 1? binom(2*i+1, i) : inv1a * (ga[i-1] + binom(2*i-1, i)*powa[i] - binom(2*i, i)*powa[i+1] + MOD*MOD) % MOD);
        gb[i] = (b == 1? binom(2*i+1, i) : inv1b * (gb[i-1] + binom(2*i-1, i)*powb[i] - binom(2*i, i)*powb[i+1] + MOD*MOD) % MOD);
        f[i] = (f[i-1] + powa[i]*gb[i] + powb[i]*ga[i] - binom(2*i, i)*powa[i]%MOD*powb[i] + MOD*MOD) % MOD;
    }

    long long res = 0;

    for (int i = 2; i <= n; i++)
        res = (res + (l[i] * binom(2*n-i-2, n-2) % MOD) * (powa[n-1] * powb[n-i] % MOD) % MOD) % MOD;
    for (int i = 2; i <= n; i++)
        res = (res + (t[i] * binom(2*n-i-2, n-2) % MOD) * (powa[n-i] * powb[n-1] % MOD) % MOD) % MOD;
    res = (res + c*f[n-2]) % MOD;

    cout << res;

    return 0;
}

문제에서 주어진 모듈러가 $10^6 + 3$으로 소수이기 때문에, 잉여역수를 쉽게 구할 수 있습니다.