AOJ2335 : 10-Year-Old Dynamic Programming / 10歳の動的計画解説(別解)

はじめに

AOJ2335:10歳の動的計画法を解いたが、非想定解で解説記事が書かれてなかったので書くことにした。別解の中で使ってるテクニックは他でも使おうと思えば使える気がします。

問題概要

2次元平面上で、次のルールで(0,0)から(N,M)まで移動するときの経路数を答えよ。

移動は複数のステップからなる
1ステップではx軸またはy軸に沿う方向に1だけ移動できる
x軸負方向またはy軸負方向への移動回数の合計はKステップである
x座標、y座標が負となるような移動はできない

制約
1 <= N,M <= 100,000, 0 <= K <= 10,000
8sec, 64MB (AOJ)

解法

想定解はJAGのページにある。この記事ではDPで解く別解について解説する。

まず、x軸方向の移動とy軸方向の移動は分けて考えることができる。

x軸上の移動を考える。k回の負方向への移動を含む移動によってx=pにたどり着いたとする。このときの経路数をf(k,p)と表記する。また、この経路数はy軸上の移動でも同じである。

x軸負方向にk回移動したのちに(N,M)にたどり着いたとしよう。すると、
x軸に沿った移動回数 N+2*k 回
y軸に沿った移動回数 M+2*(K-k) 回
となる。

つまり、答えは次の式で求まることになる。
$\sum_{k=0}^{k=K} f(k,N)f(K-k,M)\binom{N+M+2K}{N+2k}$

二項係数の計算はO(N+M+K)ぐらいの前計算からO(1)で求まるため、前計算の結果と、f(k,N)とf(k,M)がk<=Kについて分かっていればO(K)で答えが求まることが分かった。f(k,N)とf(k,M)をk<=Kについて求める方法を考えればよい。

O(K * (N+M+K))のDPをすることでf(k,N)とf(k,M)は求まるのだが、これでは時間がかかりすぎてしまう。そうであっても、行うDPは別解とそう変わらないのでこのDPについても説明しておく。

dp(i,k) := i回の移動を行い、そのうちk回が負方向への移動であるような経路数。
として、更新は次の式で行う
$dp(i,k) = \begin{cases} dp(i-1,k) + dp(i-1,k-1) & (i \geq 2k \text{ and } k \geq 1) \\ dp(i-1,k) & (\text{else}) \end{cases}$
このようにしてDPを行えば、 $f(k,p) = dp(p+2k,k)$
としてf(k,M)とf(k,N)を求めることができる。

さて、K回以上正方向に移動を行っていれば、その後どのような順番で移動しようともルールを満たしながら移動することができる。つまり、2*K回以上の移動を行った後であれば、どのタイミングで負方向への移動を行っても構わないことになる。

2*K回の移動が終わるまではDPで経路数をもとめ、2*K回以上の移動を行う経路については2*K回の移動が終わった時点のDPの結果をもとに計算すればよい。

つまり、
$dp(i,k) = \sum_{k'=max(0,2K-i+k)}^{k'=k} dp(2K,k') \binom{i-2K}{k-k'}$
でdp(i,k) (i > 2*K) を求めることができる。和をとる範囲について補足する。2*Kステップにk'回の負方向への移動が含まれていた場合、正方向への移動は2*K-k'回である。全体で正方向への移動はi-k回であるから、2*K-k' <= i-k が成立しなければならない。よって、和はk'=max(0,2K-i+k)からとる。

計算量は、DPを途中まで行うのにO(K^2)、DPの結果から、dp(N,k)とdp(M,k)を計算しf(k,N)とf(k,M)を求めるのにO(K)これが0<=k<=Kについて計算されるからO(K^2)となり、f(k,N)とf(k,M)を計算する処理全体でO(K^2)となる。

問題全体を通した計算量は、O(N+M+K^2)である。

実装

#include <iostream>
#include <algorithm>
using namespace std;
 
typedef long long ll;
 
const ll mod = 1e9+7;
ll tab[10100];
ll tmp[10100];
ll ver[10100];
ll hol[10100];
 
ll fact[300100];
ll inv_fact[300100];
 
int main() {
    fact[0] = inv_fact[0] = 1;
    for(ll i = 1; i < 300100; i++) {
        ll t = mod-2, x;
        fact[i] = (fact[i-1] * i)%mod;
        x = fact[i];
        inv_fact[i] = 1;
        while(t) {
            if(t&1) {
                inv_fact[i]*= x;
                inv_fact[i]%= mod;
            }
            x = x*x%mod;
            t >>= 1;
        }
    }
    ll N, M, K;
    cin >> N >> M >> K;
    tab[0] = 1;
    for(int i = 1; i <= 2*K ; i++) {
        for(int j = 0; j <= K; j++) {
            tmp[j] = 0;
            tmp[j] = (((i >= 2*j && j) ? tab[j-1] : 0)
                + tab[j]) % mod;
        }
        for(int j = 0; j <= K; j++) {
            tab[j] = tmp[j];
            if(i - 2*j == N) {
                ver[j] += tab[j];
                ver[j] %= mod;
            }
            if(i - 2*j == M) {
                hol[j] += tab[j];
                hol[j] %= mod;
            }
        }
    }
    for(int i = 0; i <= K; i++) {
        for(int j = 0; j <= i; j++) {
            if(2*K-N-i <= j && (N+2*i>2*K)) {
                ver[i] += (tab[j] * fact[N+2*i-2*K] % mod)
                    * (inv_fact[i-j] * inv_fact[N+2*i-2*K - (i-j)] % mod);
                ver[i] %= mod;
            }
            if(2*K-M-i <= j && (M+2*i>2*K)) {
                hol[i] += (tab[j] * fact[M+2*i-2*K] % mod)
                    * (inv_fact[i-j] * inv_fact[M+2*i-2*K - (i-j)] % mod);
                hol[i] %= mod;
            }
        }
    }
    ll res = 0;
    for(int i = 0; i <= K; i++) {
        res += ((ver[i] * hol[K-i] % mod) *
            ((fact[2*K+M+N] * inv_fact[N+2*i] % mod) * inv_fact[2*K+M+N -(N+2*i)] % mod)) % mod;
        res %= mod;
    }
    cout << res << endl;
}