mayoko’s diary

プロコンとかいろいろ。

Codeforces Round #223 (Div. 1) A. Sereja and Prefixes

解法

2 つのタイプの操作がありますが, どちらのタイプでも 10^5 番目以降の数が数の生成に関わることはないです。なのでとりあえず 10^5 個の数は配列で持っておくことにしましょう。この配列を v とします。
他に, 今まで並んでいる整数の数 len, メモ用の配列 map mp, map mpp を持っておきます。

で, 各操作では以下のようにします。

タイプ 1: v.size が 10^5 以下の時は, v に x を追加する。そうでない時は, mp[len] に x とメモしておく。
タイプ 2: v.size が 10^5 以上になるまでは, c を 1 ずつ減らしながら v に要素を正直に加えていく。size が 10^5 を超えたら, mpp[len] に l とメモして len を更新する。

すると, 10^5 番目以下のクエリには v を参照するだけで答えられ, タイプ 1 で 10^5 番目以降に来たやつは mp で答えられ, それ以外は mpp の情報から答えられます。

この問題, サンプルが弱くて上の実装を確かめるのが難しそうに思えますが, 10^5 を 10 にしてデバッグすれば, 解法がただしそうかを確かめることが出来ます。こういう時は数字ベタ書きじゃなくて ちゃんと定義したほうが嬉しいですね。

下のコードでは mpp の型が mpp となっていますがそうする必要はないです。

const int MAXSIZE = 100010;

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    int n;
    cin >> n;
    vector<int> v;
    map<ll, int> mp;
    map<ll, pii> mpp;
    ll len = 0;
    for (int i = 0; i < n; i++) {
        int type;
        cin >> type;
        if (type == 1) {
            int x;
            cin >> x;
            if (len < MAXSIZE) v.push_back(x);
            else mp[len] = x;
            len++;
        } else {
            int l, c;
            cin >> l >> c;
            while (len < MAXSIZE && c > 0) {
                for (int j = 0; j < l; j++) {
                    v.push_back(v[j]);
                    len++;
                }
                c--;
            }
            if (c) {
                mpp.insert(make_pair(len, pii(l, c)));
                len += (ll)c * l;
            }
        }
    }
    int m;
    cin >> m;
    for (int i = 0; i < m; i++) {
        ll index;
        cin >> index;
        index--;
        if (index < v.size()) cout << v[index];
        else if (mp.find(index) != mp.end()) {
            cout << mp[index];
        } else {
            auto it = mpp.upper_bound(index);
            it--;
            int l = it->second.first;
            ll tmp = it->first;
            cout << v[(index-tmp)%l] ;
        }
        cout << " ";
    }
    cout << endl;
    return 0;
}