AtCoder Regular Contest 027: D - ぴょんぴょんトレーニング

問題記録 AtCoder ARC ARC #027 平方分割 DP DP遷移の行列化

https://atcoder.jp/contests/arc027/tasks/arc027_4

問題概要

N個の石が一直線に並んでおり、石x からは 1 ～ Hx 個先のどれかの石にジャンプすることができる。以下の形式のクエリがD個飛んでくるのでそれぞれに解答せよ。

  • 整数s, t (s < t) が与えられる。石sからスタートし、ちょうど石tにたどり着くような移動経路は何通りあるか？

N <= 3 × 105

Hx <= 10

D <= 5000

解法

Hがたかだか10なので、DPの遷移を10×10の行列で表すことができる (dp[11], dp[10], ..., dp[2]) = A ・ (dp[10], dp[9], ..., dp[1]) のような形)。あとはよくあるやつで行列をセグ木とか平方分割とかに乗せればよさそうなのだが、やたら定数倍が重くメモリ制限も小さめなので適当にやるとだめになる (3*105 * 102 * 64 bits で既に240MB)。自分は平方分割でやったが、合成前の個別の行列については陽に持たず必要になるたびいちいち生成するようにした。あとは行列の積を取る方向が微妙にややこしい（大きいインデックスのほうが左に来る）のでがんばる。あと公式解説にあるとおり各クエリに答えるときは行列同士をかけるのではなくベクトルにかけていくようにして時間計算量を抑えた。

感想

虚無みがあったが、DPの遷移を行列でやるやつについての理解は深まったのでよかった

コード (C++)

#include <bits/stdc++.h>
using namespace std;
#define REP(i,n) for (int i=0;i<(n);i++)
#define REP2(i,m,n) for (int i=m;i<(n);i++)
typedef long long ll;

const ll MOD = 1000000007;
const int SEG = 600;
const int M = 10;

int N, Q, s, t;
int H[303030];
ll bucket[SEG][M][M];
ll tmp[M][M];
ll tmp2[M][M];
ll vec[M];
ll vec_tmp[M];

void init_tmp() {
    REP(i, M) tmp[i][i] = 1;
}

void init_vec() {
    vec[0] = 1;
    REP2(i, 1, M) vec[i] = 0;
}

void gen(int i) {
    REP(j, M) REP(k, M) tmp[j][k] = 0;
    REP2(j, 1, M) tmp[j][j-1] = 1;
    REP2(j, 1, M+1) if (i - j >= 0 && H[i-j] >= j) tmp[0][j-1] = 1;
}

void prod1(int i) {
    gen(i);
    REP(j, M) vec_tmp[j] = vec[j];
    REP(j, M) vec[j] = 0;
    REP(j, M) REP(k, M) (vec[j] += tmp[j][k] * vec_tmp[k] % MOD) %= MOD;
}

void prod2(int seg) {
    REP(i, M) vec_tmp[i] = vec[i];
    REP(i, M) vec[i] = 0;
    REP(j, M) REP(k, M) (vec[j] += bucket[seg][j][k] * vec_tmp[k] % MOD) %= MOD;
}

void solve() {
    cin >> N;
    REP(i, N) cin >> H[i];
    cin >> Q;

    REP(i, SEG) REP(j, M) bucket[i][j][j] = 1;
    for (int i = 0; i < N; ++i) {
        gen(i);
        REP(x, M) REP(y, M) tmp2[x][y] = bucket[i/SEG][x][y];
        REP(x, M) REP(y, M) bucket[i/SEG][x][y] = 0;
        REP(x, M) REP(y, M) REP(z, M) (bucket[i/SEG][x][y] += tmp[x][z] * tmp2[z][y]) %= MOD;
    }

    while (Q--) {
        cin >> s >> t;
        --t;
        init_vec();
        for (int i = s / SEG * SEG; i <= t / SEG * SEG; i += SEG) {
            if (i < s || t < i + SEG - 1) {
                for (int j = max(i, s); j <= min(t, i + SEG - 1); ++j) {
                    prod1(j);
                }
            } else {
                prod2(i / SEG);
            }
        }
        cout << vec[0] << "\n";
    }
}

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    solve();
}