2020-11-15

AtCoder ABC 170 E - Smart Infants (水色, 500 点)

こういうの苦手すぎるので、素早く綺麗に実装できるようになりたい

問題へのリンク

問題概要

AtCoder に参加している幼児が $N$ 人おり、 $1$ から $N$ の番号が付けられている。また、幼稚園 $1, 2, \dots, 2 \times 10^{5}$ がある。幼児 $i$ のレートは $A_{i}$ であり、はじめ幼稚園 $B_{i}$ に所属している。

これから $Q$ 回にわたって、転園が行われる。 $j$ 回目の転園では、幼児 $C_{j}$ の所属を幼稚園 $D_{j}$ に変更する。

$Q$ 回それぞれの転園後における、「AtCoder に参加している幼児が一人以上いる幼稚園それぞれにおける園内でのレーティング最高値の最小値」を求めよ。

制約

$1 \le N, Q \le 2 \times 10^{5}$

考えたこと

まず、以下の処理を高速に実現できる必要がある！！

【削除】幼稚園 $d$ において、幼児 $i$ (レーティング $A_{i}$ ) が抜けることによって、幼児園 $d$ 内の最高レーティングがどのように変動するか
【挿入】幼稚園 $d$ において、幼児 $i$ (レーティング $A_{i}$ ) が入ってくることによって、幼稚園 $d$ 内の最高レーティングがどのように変動するか

これらは幼稚園ごとに以下のようなデータ構造を持つことで実現できる。

king[ d ] := 幼稚園 $d$ において、所属する各幼稚 $a$ についての (幼稚 $a$ のレーティング, $a$ ) のペアを格納したもの (たとえば set<pair<int,int>> 型)

このようにすることで、幼稚園からの幼稚の削除・挿入が高速にできるだけでなく、幼稚の最高レーティングもわかる！ちなみに、set に「幼稚のレーティング」だけでなく、「幼稚の id」もペアにして格納しているのは、レーティングが重複する場合に対処するため。

各幼稚園の最高レーティングの最小値

以上より、幼稚 $a$ が幼稚園 $s$ から幼稚園 $t$ へと移動することに伴う

幼稚園 $s$ の最高レーティングの変化
幼稚園 $t$ の最高レーティングの変化

は高速に扱えるようになった。今度は幼稚園全体についての情報を管理するために以下のクエリを処理するデータ構造が必要となる。

【変更】幼稚園 $d$ における最高レーティングが $x$ から $y$ へと変更する
【取得】すべての幼稚園の最高レーティングの最小値を求める (幼稚のいない幼稚園の最高レーティングは INF とする)

これを実現する方法はいくつか考えられる。

セグメント木を用いて RMQ を実装する (割と楽だと思う)
set を用いて管理する (想定解法)

などがある。ここではセグメント木を用いてみることにした。このとき、幼稚園の個数を $M$ として、計算量は $O(N \log N + M + Q \log M)$ となる。

コード

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

template<class Monoid> struct RMQ {
    const Monoid INF;
    int SIZE_R;
    vector<pair<Monoid,int> > dat;
    
    RMQ(int n, const Monoid &inf): INF(inf) { init(n); }
    void init(int n) {
        SIZE_R = 1;
        while (SIZE_R < n) SIZE_R *= 2;
        dat.assign(SIZE_R * 2, pair<Monoid,int>(INF, -1));
    }
    
    /* set, a is 0-indexed */
    void set(int a, const Monoid &v) { dat[a + SIZE_R] = make_pair(v, a); }
    void build() {
        for (int k = SIZE_R - 1; k > 0; --k) {
            dat[k] = min(dat[k*2], dat[k*2+1]);
        }
    }
    
    /* update, a is 0-indexed */
    void update(int a, const Monoid &v) {
        int k = a + SIZE_R;
        dat[k] = make_pair(v, a);
        while (k >>= 1) dat[k] = min(dat[k*2], dat[k*2+1]);
    }
    
    /* get {min-value, min-index}, a and b are 0-indexed */
    pair<Monoid,int> get(int a, int b) {
        pair<Monoid,int> vleft = make_pair(INF, -1), vright = make_pair(INF, -1);
        for (int left = a + SIZE_R, right = b + SIZE_R; left < right; left >>= 1, right >>= 1) {
            if (left & 1) vleft = min(vleft, dat[left++]);
            if (right & 1) vright = min(dat[--right], vright);
        }
        return min(vleft, vright);
    }
    inline Monoid operator [] (int a) { return dat[a + SIZE_R].first; }
    
    /* debug */
    void print() {
        for (int i = 0; i < SIZE_R; ++i) {
            Monoid val = (*this)[i];
            if (val < INF) cout << val;
            //else cout << "INF";
            if (i != SIZE_R-1) cout << ",";
        }
        cout << endl;
    }
};

using pll = pair<long long, long long>;
const int MAX = 210000;
const long long INF = 1LL<<60;

int main() {
    int N, Q;
    cin >> N >> Q;
    vector<long long> A(N), B(N, -1);
    RMQ<long long> rmq(MAX+1, INF);
    vector<set<pll>> king(MAX);

    auto move = [&](int s, int from, int to) -> long long {
        if (from != -1) {
            king[from].erase(pll(-A[s], s));
            long long val = (!king[from].empty() ? -king[from].begin()->first : INF);
            rmq.update(from, val);
        }
        B[s] = to;
        king[B[s]].insert(pll(-A[s], s));
        rmq.update(B[s], -king[B[s]].begin()->first);
        return rmq.get(0, MAX+1).first;
    };

    for (int i = 0; i < N; ++i) {
        cin >> A[i] >> B[i];
        --B[i];
        move(i, -1, B[i]);
    }
    for (int q = 0; q < Q; ++q) {
        int c, d;
        cin >> c >> d;
        --c, --d;
        cout << move(c, B[c], d) << endl;
    }
}

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

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

AtCoder ABC 170 E - Smart Infants (水色, 500 点)

問題概要

制約

考えたこと

各幼稚園の最高レーティングの最小値

コード