けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder AGC 003 E - Sequential operations on Sequence (赤色, 1400 点)

この頃、数列を繰り返すのが流行ってたのかな

問題へのリンク

問題概要

長さ  N の恒等数列 ( 1, 2, \dots, N) が与えられる。この数列に以下の操作を合計で  Q 回行う。 i 番目の操作は、パラメータ  q_{i} であらわされ、以下のように行われる。

  • 現在の数列を無限回繰り返した数列の先頭  q_{i} 項をとって、新たな数列とする。

 Q 回の操作後、この数列に  1 から  N までの各々の数が何回ずつ現れるかを求めよ。

制約

  •  1 \le N, Q \le 10^{5}
  •  1 \le q_{i} \le 10^{18}

考えたこと

ひとまず思ったことは

  • もし  N より小さい  q が来たら、それまでの経緯がどのようなものであったとしても、操作後の数列は  1, 2, \dots, q の繰り返しになる
  •  q \gt r であるとき、 q による操作は無視しても  r による操作後の結果は変わらない

といったあたり。よって、特に二番目の考察から、 q は狭義単調増加である場合に帰着される。...で、この問題は AGC なんだから...というメタ読みにしたがうと、きっと変なデータ構造は使わずに解けるに違いない。あと簡単のため、便宜上、

  •  q_{0} \gt N だったら、 q の先頭に  N を付け加えておく
  •  q_{0} \lt N だったら、 N = q_{0} としておく

ということにしておく。

 N = 3,  q = (3, 5, 14, 22) とかでやると、こんな感じになる。

1, 2, 3
1, 2, 3, 1, 2
1, 2, 3, 1, 2, 1, 2, 3, 1, 2, 1, 2, 3, 1
1, 2, 3, 1, 2, 1, 2, 3, 1, 2, 1, 2, 3, 1, 1, 2, 3, 1, 2, 1, 2, 3

(1 ... i) ごとにまとめると、この変化は

  • (1 ... 3) が 1 個
  • (1 ... 3) が 1 個、(1 ... 3) が 1 個
  • (1 ... 3) が 3 個、(1 ... 2) が 2 個、(1 ... 1) が 1 個
  • (1 ... 3) が 5 個、(1 ... 2) が 3 個、(1 ... 1) が 1 個

という感じになっている。この「係数」を上手に求めたい。

breakdown して考える

たとえば  N = 3 q = (3, 5, 7, 23, 30, 78) とかだったら、

  • 78 = 30 × 2 + 18

とかになる。この 18 の部分がどうなっているのかが問題だけど、これは結局、 q = 23 のクエリに答えた時代に形作られたもののうち、先頭の 18 個が並んでいるものと一致することがわかる。よってその状況は、数列の長さが 7 である状態を繰り返して 18 にするのに等しい。

  • 18 = 7 × 2 + 4

となり、

  • 4 = 3 × 1 + 1

となって、操作は完結する。まとめると

  • 78 = 30 × 2 + 7 × 2 + 3 × 1 + 1

というわけだ。こうして、 q = 78 のケースを、その手前の  q = (3, 5, 7, 23, 30) の線形結合で表すことができた。よって、もし  q = 3, 5, 7, 23, 30 のそれぞれの場合について、(1 ... 3) と (1 ... 2) と (1 ... 1) の係数がわかっていたら、 q = 78 の場合の係数も求められることになる。

  • 1: (0, 0, 1)
  • 2: (0, 1, 0)
  • 3: (1, 0, 0)
  • 5: (1, 1, 0)
  • 7: (1, 2, 0)
  • 23: (3, 7, 0)
  • 30: (4, 9, 0)
  • 78: (11, 22, 1)

というわけで、答えは 1 が 34 個、2 が 33 個、3 が 11 個である。

後ろから経路数

このままだと、空間計算量が  O(NQ) を要してしまう。しかしよく考えると、

  • 78 = 30 × 2 + 7 × 2 + 3 × 1 + 1

とか表すときの、線形結合の係数だけ気にすればよさそうだ。それでなんとかなりそうに思える。具体的には、次のようなグラフにおいて、私たちは

  • ノード 3 からノード 78 への経路数
  • ノード 2 からノード 78 への経路数
  • ノード 1 からノード 78 への経路数

をボトムアップに求めているような感じになっている。でも逆に、ノード 78 からトップダウンに経路数を求めても良い。そうすると、このグラフの枝数分の計算量で済む。

計算量としては、各ノードから出る矢印の本数は、たとえば 78 の変化に着目すると

78 → 18 → 4 → 1

という具合に、必ず半減以下になることがわかるので、枝の本数は  O(Q\log{A}) ( A はクエリの値の最大値) である。ただし枝を引く場所を求めるのに二分探索を使うので、計算量としては  O(Q\log{Q}\log{A} + N) となる。

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

int N, Q;
deque<long long> iq, q;

vector<long long> solve() {
    if (Q == 0) return vector<long long>(N, 1);
    
    // stack
    int iter = 0;
    for (int i = 0; i < Q; ++i) {
        while (iter && q[iter-1] >= iq[i]) --iter;
        q[iter++] = iq[i];
    }
    q.resize(iter);
    int NN = N;
    if (N >= q[0]) NN = q[0];
    else q.push_front(N);

    // dp
    map<long long, long long> dp;
    dp[q.back()] = 1;
    for (int i = q.size()-1; i > 0; --i) {
        long long cur = q[i], nex = q[i-1];
        dp[nex] += dp[q[i]] * (cur / nex);
        cur %= nex;
        while (cur > NN) {
            int id = upper_bound(q.begin(), q.end(), cur) - q.begin() - 1;
            nex = q[id];
            dp[nex] += dp[q[i]] * (cur / nex);
            cur %= nex;
        }
        if (cur) dp[cur] += dp[q[i]];
    }
    vector<long long> res(N, 0);
    res[NN-1] = dp[NN];
    for (int i = NN; i > 1; --i) res[i-2] = res[i-1] + dp[i-1];
    return res;
}

int main() {
    cin >> N >> Q;
    iq.resize(Q); q.resize(Q);
    for (int i = 0; i < Q; ++i) cin >> iq[i];
    const auto &res = solve();
    for (int i = 0; i < N; ++i) cout << res[i] << endl;
}