【ABC179】 D - Leaping Tak

atcoder.jp

解法

ぱっと見、部分和問題かなと思うが、O(N²)になってしまうため間に合わない。

他の方法を考えてみると、 ・なぜか区間が与えられている ・Kが小さい ・区間は共通部分を持たない と怪しい部分がたくさんあることに気づく。
区間を利用して高速にDPができればよいので、セグメント木が使えそう。
区間加算(RAQ)・一点取得さえできればよいので遅延評価である必要はない。

DPの定義

dp[i] ← iまで行く場合の数

初期値: dp[1] = 1; dp[i] = 0; (1 < i ≦ N)
遷移: dp[i + d] += dp[i] (0 ≦ k < K, L[k] ≦ d ≦ R[k])

遷移を行う際にセグメント木が強力。

実装

struct RAQ_SegTree {
    int n; // 葉の数
    vector<mint> node; // 完全2分木

    // 初期化
    RAQ_SegTree(vector<long long> v) {
        int sz = v.size();
        n = 1;
        while (n < sz) n *= 2; // 与えられた数列の項数以上の2^n個、葉を作る
        node.resize(n*2-1, 0); // 0で初期化

        for (int i = 0; i < sz; i++) node[n-1+i] = v[i]; // 葉を初期化
        // 下から順に葉以外のnodeを初期化
        for (int i = n-2; i >= 0; i--) node[i] = 0;
    }

    // 取得
    mint get(int i) { // iは数列の添字(0-index)
        mint ans = 0;
        i += n-1; // 葉はn-1から始まる
        ans += node[i]; // 葉の値を加算
        while (i > 0) { // 親の値を加算
            i = (i-1)/2; // 親の添字
            ans += node[i];
        }
        return ans;
    }

    // [s, t)にxを加算
    void update(int s, int t, mint x) { update_query(s, t, 0, n, 0, x); }

    void update_query(int s, int t, int l, int r, int n, mint x) {
        if (r <= s || t <= l) return; // 範囲外なら終了
        // [s, t)が[l, r)を内包しているとき
        else if (s <= l && t >= r) node[n] += x;
        else {
            // (r+l)/2は区間の中心, 区間の中心を左端にするか、右端にするかで分岐する
            update_query(s, t, l, (r+l)/2, n*2+1, x); // 左下の子を更新
            update_query(s, t, (r+l)/2, r, n*2+2, x);  // 右下の子を更新
        }
    }

    // デバック用
    void output() {
        for (int i = 0; i < n*2-1; i++) cout << node[i] << " ";
    }

};

void solve(long long N, long long K, std::vector<long long> L, std::vector<long long> R){
    vector<long long> seg(N+1, 0);
    seg[1] = 1;
    auto segtree = RAQ_SegTree(seg);

    for (int i = 1; i < N; i++) {
        mint p = segtree.get(i);
        REP (j, K) {
            if (i + L[j] <= N) {
                segtree.update(i + L[j], min(N, i + R[j])+1, p);
            }
        }
    }

    c(segtree.get(N))
}