[AtCoder] E – BBQ Hard(AGC 001)

スポンサーリンク

概要

問題文 →
自分の提出 →
公式解説 →

経路数に言い換える頭の良さそうな解法しか見なかったため、書きます。
二項係数について何か上手くやりたいときは、たいていの場合は、計算対象を多項式についての何かだと読み替えると、ごく自然に解に至ることができます(効用には個人差があります)。

二項係数 ${}_n\mathrm{C}_r$ を $\binom{n}{r}$ と表します。また、$A_i + B_i = C_i$ と表すことにします。また、$D = 2000$ とします($A_i, B_i$ の上界)。

解法

この問題は、立式までは超簡単です。

問題:$\sum_{1\leq i < j\leq N}\binom{C_i+C_j}{A_i+A_j}$ を求めよ.

二項係数といえば二項定理です。二項定理により言い換えると \[ \binom{C_i+C_j}{A_i+A_j} = (1+x)^{C_i+C_j}のx^{A_i+A_j} の係数\] となります。

うまく $i$ や $j$ についてまとめて計算できそうな形にしたいです。単に $\sum_{i,j} (1+x)^{C_i+C_j}$ であれば簡単なのですが、集計対象となる次数が異なるところが厄介です。そこで、\[ \binom{C_i+C_j}{A_i+A_j} = x^{-A_i-A_j}(1+x)^{C_i+C_j} の定数項 \] と読み替えてみます(定数項の和は、和の定数項なので、この方が扱いやすい)。

さらに、$f_i = x^{-A_i}(1+x)^{C_i}$ と書くことにすると、\[ \binom{C_i+C_j}{A_i+A_j} = f_if_j の定数項 \] と言い換えられました。これで、和がまとめられそうな形になります。

問題の言い換え:$\sum_{1\leq i < j\leq N} f_if_j$ の定数項を求めよ。

ここで、一般的なテクニック $\sum_{1\leq i < j\leq N}a_ia_j = \frac12\bigl((a_1+\cdots+a_n)^2 – (a_1^2 + \cdots + a_n^2)\bigr)$ を利用します。

・$f = \sum_i f_i$ とする。$f^2$ の定数項を $X$ とする。
・$\sum_{i} f_i^2$ の定数項を $Y$ とする。
・$(X-Y)/2 \pmod{p}$ が答である。

$Y$ の計算は簡単です。$f_if_j$ の定数項が二項係数でしたから、$Y = \sum_i \binom{C_i+C_i}{A_i+A_i}$ となり、階乗の事前計算のもと 1件あたり $\Theta(1)$ で計算できます。よって、残ったものは $f$ の計算です。

残った課題:$f = \sum_i f_i = \sum_{i} x^{-A_i}(1+x)^{C_i}$ を計算せよ。

これを $i$ ごとにばらばらに計算するのでは、計算が厳しいです。$D$ 次式を $N$ 個足すことになってしまいます。何とかして、$x^{-A_i}$ に対する固有の処理を1回で済ませられる形にします。

$f$ を、$(1+x)$ の指数ごとにまとめて $f = \sum_{C} f_C (1+x)^C$ などと書いてみます。各 $f_C$ に対して $(1+x)^C$ をかけてから集計するのが無駄で、次の要領で計算を行うと、各 $f_C$ をそのままの形で1度ずつ触るだけで済みます:

\[f_1(1+x) + f_2(1+x)^2 + f_3(1+x)^3 = (1+x)\bigl(f_1 + (1+x)\bigl(f_2 + (1+x)f_3\bigr)\bigr)\]

つまり、例えば $n = \max C$ から始めて、$n$ を小さくしながら「$(1+x)$倍 → $f_n$ を加える」を適切なところまで行えばよいです。$10$ 進法表記された巨大な数を調べるときなどにもある計算方法ですね。

計算過程で現れる式はすべて $x^{-D},\ldots,x^D$ の線形結合で、$(1+x)$ 倍操作は1件あたり $\Theta(D)$ で行えます。$f$ の計算は、全体として計算量は$\Theta(D^2 + N)$ で行えます。あとは、$f^2$ の定数項 $X$ を求めて $(X-Y)/2$ を答えればよいです。

実装用メモ

道中で出てくるのが、多項式ではなく負べきの項も含む (環 $\mathbb{Z}[x,1/x]$ の元)のがちょっと邪魔です。全て $x^D$ をかけておくと、登場人物が全て $2D$ 次以下の多項式になります。

・$f = 0$ から始める。$n = 2D, 2D-1, \ldots, 0$ に対して次を行う。
・$f$ を $(1+x)$ 倍したあと、$a+b = n$ となる $(a,b)$ に対して $x^{D-a}$ を加える。

「$f^2$ の定数項」としていたものは、この場合には $2D$ 次の係数に置き換わり、$2D+1$ ペアの積の総和をとるだけです。

numpy の場合

C_to_A には、$c = a+b$ に対応する $D-a$ のリストが入っています。

f = np.zeros(D+D+1,np.int64)
for A in C_to_A[::-1]:
f[1:] += f[:-1].copy()
np.add.at(f,A,1)
f %= MOD
X = (f * f[::-1] % MOD).sum() % MOD

.copy() は、numpyがよほど古くなければ不要。

タイトルとURLをコピーしました