AOJ 2606 Perm Query

問題文

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2606

解法

まず，制約を読むと，ある項についてたかだか40回の遷移で元に戻ることがわかるので，愚直にこれを求めます．
このとき，それぞれの項が c[i] 回 で元に戻るとします．
すると，[l, r] が与えられた時，この区間がちょうど最初の状態に戻るのは，L = lcm(c[l], c[l+1], ..., c[r]) 回後であることは明らかです．
そして，各項は L / c[i] 回サイクルを回ることになります．1サイクルでの和s[i] を求めておけば，これは s[i] * L / c[i] を l から r まで足した和が答えになります．

区間の LCM は segment tree で求められるのでそうします．

このままだと O(QNlogN) なのですが，sum[k][i] := （第 i 項までの，サイクルが k であるようなものの s[i] の和）という累積和を持っておけば，k が L の約数であるようなものについてのみ (L / k) * (sum[k][r] - sum[k][l]) を考えて総和を取れば，求めたい答えになるのでそうします．k はたかだか 40 なので間に合います．

ソースコード

#include <bits/stdc++.h>
using namespace std;

using ll = long long;

constexpr ll M = 1e9 + 7;
constexpr int MAX_K = 41;

// セグツリは略

int main() {
    int N, Q;
    cin >> N >> Q;
    vector<ll> p(N);
    for(int i = 0; i < N; ++i) {
        cin >> p[i];
    }

    vector<ll> cycle(N);
    vector<vector<ll>> sum(MAX_K, vector<ll>(N + 1));
    for(int i = 0; i < N; ++i) {
        ll s = p[i];
        ll now = p[i];
        int k = 1;
        for(k = 1; k < MAX_K; ++k) {
            if(now != i + 1) {
                now = p[now - 1];
                s = (s + now) % M;
            } else {
                break;
            }
        }
        cycle[i] = k;
        sum[k][i + 1] = s;
        for(int k = 0; k < MAX_K; ++k) {
            sum[k][i + 1] = (sum[k][i] + sum[k][i + 1]) % M;
        }
    }
    segment_tree<LCM> lcm_seg(cycle);
    while(Q--) {
        int l, r;
        cin >> l >> r;
        l--;
        ll L = lcm_seg.query(l, r);
        ll res = 0;
        for(int k = 1; k < MAX_K; ++k) {
            if(L % k == 0) {
                (res += ((L / k) % M) * (sum[k][r] - sum[k][l] + M)) %= M;
            }
        }
        cout << res << endl;
    }
}