つちのこの日記

競プロと音ゲーがメインです。

キーエンス プログラミング コンテスト 2019 - F Paper Cutting(赤diff, 900点) 形式的冪級数解

問題リンク

atcoder.jp

 

問題概要

縦に $H$ 本、横に $W$ 本の線が引かれた長方形の紙がある。

一回の操作では、まだ選択してない線を一つ選択し、その線で紙を切る。その後紙の破片の個数をスコアに加算する。

操作をちょうど $K$ 回行う時、全ての操作方法に対するスコアの総和を求めよ。

制約
  • $1 \le H, W \le 10^7$
  • $1 \le K \le H + W$

解説

本解説では $[x^n]f(x)$ で形式的冪級数 $f$ の $n$ 次の項の係数を表すこととする。

$i$ 回目の操作で縦の線を選択した時のスコアの増加分の期待値を求めることを考える。これを $\mathcal{Ο}(1)$ で求められれば、この問題に $\mathcal{Ο}(K)$で答えられる。

考察を容易にするために、操作は 0_indexedで考える(つまり、操作は $0$ 回目から $K - 1$ 回目まで行う)。

$i$ 回目の操作の前に既に選択された縦線の個数を $j$ 本と固定してみる。この時、既に選択された横線の個数は $i - j$ 本である。

このような状態になっている確率は、 $ \frac{\displaystyle \binom {H-1}{j} \displaystyle \binom{W}{i-j}}{\displaystyle \binom{H+W-1}{i}} $ であり、この時得られるスコアは $ ( j + 2 )( i - j + 1) $ である。

ゆえに、 $i$ 回目の操作で縦の線を選択した時のスコアの増加分の期待値は以下の数式で表せる。 $$\displaystyle \sum_{j = 0} ^ i \left( \frac{\displaystyle \binom {H-1}{j} \displaystyle \binom{W}{i-j}}{\displaystyle \binom{H+W-1}{i}}( j + 2 )( i - j + 1) \right) \tag{*}$$ 

分母は $j$ に依存せず楽に計算できるので、後は分子を上手く扱う。

分子を $\binom {H-1}{j}(j + 2) $ と $\binom{W}{i-j}(i-j+1)$ に分けて考える。

一つ目については、 $$f(x) = \sum_{i = 0}^{\infty}\binom{H-1}{i} x^i$$ と置けば、 $$xf'(x) = \sum_{i = 0}^{\infty}{i \binom{H-1}{i} x^i}$$ であるから、

$$\begin{split} \binom {H-1}{j}(j + 2) & = [x^j] 2f(x) + xf'(x) \cr & = [x^j]( (H+1)x + 2)(1+x)^{H-2} \cr\end{split} $$ と表すことができる。

同様に、二つ目については $u := i-j$ と置けば

$$\begin{split} \binom{W}{i-j}(i-j+1) & = \binom{W}{u}(u+1)) \cr & = [x^u]( (W+1)x + 1)(1+x)^{W-1} \end{split} $$ と表せる。

ゆえに $${(*)} = \frac{1} {\binom{H+W-1}{i}}\displaystyle \sum_{j = 0} ^ i \left( [x^j] (H+1)x + 2)(1+x)^{H-2} × [x^{i-j}]( (W+1)x + 1)(1+x)^{W-1} \right)$$

これは畳み込みの形となっているので、求めるものは $$\frac{1} {\binom{H+W-1}{i}} [x^i] (1+x)^{H+W-3}( (H+1)x + 2) ( (W+1)x + 1) $$ と簡潔に表示できた。

$(1+x)^N$ の 任意の項の係数は、適切に階乗とその逆元を求めることで $\mathcal{Ο}(1)$ で求められるので、$i$ 回目の操作で縦の線を選択した時のスコアの増加分の期待値も $\mathcal{Ο}(1)$ で求められる。

縦線が選ばれる確率は $\frac{H}{H+W}$ 、横線が選ばれる確率は $\frac{W} {H+W}$ であるから、先ほどの期待値にこの確率を掛け合わせ、$i = 0$ から $K - 1$ まで足し合わせることでこの問題を $\mathcal{Ο}(K)$ で解くことができた。

実装の際には階乗と逆元の計算を適切に行うことや、負の二項係数にライブラリが対応していない場合 $H=W=1$ の時に正しい答えが返されないことに注意されたい。

ソースコード 

 

長いので折り畳み(includeやmodintは省略)
const int CombMAX = 20000010;
mint fac[CombMAX + 1], finv[CombMAX + 1], inv[CombMAX + 1];

struct Combinationinit {
    Combinationinit() {
        fac[0] = fac[1] = 1;
        finv[0] = finv[1] = 1;
        for(int i = 2; i <= CombMAX; i++) {
            fac[i] = fac[i - 1] * i;
        }
        finv[CombMAX] = fac[CombMAX].pow(MOD - 2);
        for(int i = CombMAX; --i;) {
            finv[i] = finv[i + 1] * (i + 1);
        }
    }
} Combinationinit_;

//nCk
mint COM(int n, int k) {
    if(n < 0) return COM(-n, k) * (k % 2 ? -1 : 1);
    if(n < k or k < 0) return 0;
    return fac[n] * finv[k] * finv[n - k];
}

mint COMinv(int n, int k) {
    if(n < 0) return COMinv(-n, k) * (k % 2 ? -1 : 1);
    if(n < k or n < 0 or k < 0) return 0;
    return finv[n] * fac[k] * fac[n - k];
}

//nPk
mint PER(int n, int k) {
    if(n < k or n < 0 or k < 0) return 0;
    return fac[n] * finv[n - k];
}

int h, w, k, i;
mint res;
int main() {
    scanf("%d%d%d", &h, &w, &k);

    for(i = 0; i < k; i++) {
        mint tres = COM(h + w - 3, i) * 2
                  + COM(h + w - 3, i - 1) * (h + 2 * w + 3)
                  + COM(h + w - 3, i - 2) * (h + 1) * (w + 1);
        mint yres = COM(h + w - 3, i) * 2
                  + COM(h + w - 3, i - 1) * (w + 2 * h + 3)
                  + COM(h + w - 3, i - 2) * (h + 1) * (w + 1);
        res += (tres * h + yres * w) * COMinv(h + w - 1, i);
    }

    cout << res * PER(h + w, k) / (h + w) << endl;

    return 0;
}