Writer 解説: https://yukicoder.me/problems/no/801/editorial

問題概要

ある$ N (\leq 3000)$階建てビルに$ M (\leq 3000) $台のエレベーターがあり，それぞれ$ L_i $階から$ R_i $階まで各階に移動することができる．

初め 1 階にいるとして，$ K (\leq 3000) $回の移動の後に$ N $階にいるような移動の仕方の通り数を$ 10^9+7 $で割ったものを求めよ．

解法概要

dp[何回移動したか][最後に何階にいるか] の動的計画法で解くことができるが，そのまま素直に実装しようとすると (例えば，i階からj階まで移動するのに使えるエレベーターの数a[i][j]のような配列を持っていても) $ O(KN^2) $ となってしまい TLE する．

そこで，エレベーターごとに$ L_i $，$ R_i $を保持しておき，エレベーターごとに dp[現在の移動回数][$ L_i $〜$ R_i $] の和を dp[現在の移動回数+1][$ L_i $〜$ R_i $] にそれぞれ足しこむことを考える．なぜならば，現在$ L_i $〜$ R_i $階のどこかにいるなら，エレベーターiを使って$ L_i $〜$ R_i $階のどこでも好きな階に動くことができるためである．

dp[現在の移動回数+1][$ L_i $〜$ R_i $] にそれぞれ足す処理は，いもす法を用いて計算量を $ O(MN^2) $ から $ O(MN) $ へと落とすことができる．しかし，これでもまだ$ O(KMN) $で，TLEしてしまう．

そこでさらに，あらかじめdp[現在の移動回数][j]についての累積和を取っておくことによってdp[現在の移動回数][$ L_i $〜$ R_i $] の和を計算する処理を高速化する．

累積和を取る処理が$ O(N) $で，そこから和を求める1回の処理が$ O(1) $なので，これで全体の計算量は $ O(K(M+N)) $ となり，時間内に処理を終えることができる．

ソースコード

#include <iostream>
#include <vector>
#include <algorithm>
#include <utility>

using namespace std;

typedef long long ll;
typedef pair<int, int> Pii;

const ll MOD = 1000000007;

int main() {
    int N,M,K;
    cin >> N >> M >> K;

    vector<Pii> elevators(M);
    for(int i=0;i<M;++i) {
        cin >> elevators[i].first >> elevators[i].second;
    }

    vector< vector<ll> > dp(K+1, vector<ll>(N+2, 0LL));
    dp[0][1] = 1LL;
    for(int i=0;i<K;++i) {
        for(int j=1;j<=N+1;++j) {
            dp[i][j] += dp[i][j-1];
            dp[i][j] %= MOD;
        }
    
        for(int j=0;j<M;++j) {
            if(!dp[i][elevators[j].first]) continue;
            ll total = (dp[i][elevators[j].second] - dp[i][elevators[j].first-1] + MOD) % MOD;
            dp[i+1][elevators[j].first] += total;
            dp[i+1][elevators[j].first] %= MOD;
            dp[i+1][elevators[j].second+1] += MOD - total;
            dp[i+1][elevators[j].second+1] %= MOD;
        }

        for(int j=1;j<=N+1;++j) {
            dp[i+1][j] += dp[i+1][j-1];
            dp[i+1][j] %= MOD;
        }
    }

    cout << dp[K][N] << endl;
}

所感

これ面白い．シンプルな問題ながら累積和と imos 法で計算量がゴリゴリ落ちるのが爽快．

コンテスト中は行列累乗だと思っていて，必死に$ O(N^3 \log K) $解を書いてしまっていた．

ところで，ブログに yukicoder の問題の解説を書くのは初めてです．

NoiminのNoise

yukicoder No.801 エレベーター

問題概要

解法概要

ソースコード

所感