【AtCoder】ABC 410 C - Rotatable Array | 緑コーダーが解くAtCoder

atcoder.jp

配点: 300 点 / 実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 230 / NoviSteps: 3Q

問題概要
制約
考察
実装例
コメント

問題概要

長さ $N$ の整数列 $A$ があり、最初 $A_i = i$ である。以下のクエリを全部で $Q$ 個処理せよ。

1 p x : $A_p$ を $x$ に変更する。
2 p : $A_p$ を出力する。
3 k : $A$ 全体を $k$ 回左巡回シフトする。

制約

入力は全て整数
$1 \le N \le 10^6$
$1 \le Q \le 3 \times 10^5$
全てのクエリは以下の制約を満たす。
- $1 \le p \le N$
- $1 \le x \le 10^6$
- $1 \le k \le 10^9$

考察

各クエリを愚直に行おうとすると、クエリ3に $O(N)$ の計算量がかかってしまい、 $Q$ 個のクエリを処理する中で TLE する。

そこで、実際に $A$ の各要素をシフトさせる代わりに、シフトが行われた回数を記録しておくことにする。

現時点でシフトが行われた回数を $s$ とするとき、 $A$ の $p$ 番目の要素は実際には $(p + s) \mod N$ 番目の位置にあることになる。ただし、これは 0-indexed である。

したがって、クエリ3が呼ばれるごとに $s$ に $k$ を加えることで、このクエリは $O(1)$ の計算量で処理することができる。

クエリ1, 2はそのまま $O(1)$ で処理できるので、全体としては $O(N + Q)$ で解くことができた。

実装例

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

#define all(x) (x).begin(), (x).end()
#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)

using ll = long long;

// ======================================== //

int main()
{
    int N, Q;
    cin >> N >> Q;

    vector<int> A(N);
    iota(all(A), 1);

    ll left_shift = 0;

    while (Q--)
    {
        int t;
        cin >> t;

        if (t == 1)
        {
            int p, x;
            cin >> p >> x;
            p--;

            A[(p + left_shift) % N] = x;
        }
        else if (t == 2)
        {
            int p;
            cin >> p;
            p--;

            cout << A[(p + left_shift) % N] << endl;
        }
        else if (t == 3)
        {
            int k;
            cin >> k;

            left_shift += k;
        }
    }

    return 0;
}