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

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

AtCoder ABC 171 D - Replacing (400 点)

差分更新を上手にやっていくという、とても教育的な問題!!!
そして、よく似た類題として、以下の問題がある!!

atcoder.jp

今回の問題へのリンク

問題概要

 N 個の整数  A_{1}, A_{2}, \dots, A_{N} が与えられる。以下の  Q 回のクエリに答えよ。

  • 各クエリでは 2 つの整数  B, C が与えられる
  • 数列の中の値  B の部分を、すべて値  C に置き換える
  • こうして得られる数列の値の総和を求めよ

制約

  •  1 \le N, Q, A, B, C \le 10^{5}

考えたこと: バケットに変換

この手の問題では、数列をそのまま扱うのではなく、バケットで管理するのがめちゃくちゃ有効!!!

つまり、数列  A_{1}, A_{2}, \dots, A_{N} を、次のような配列に変換して考える。

  • num[ v ] := 数列  A_{1}, A_{2}, \dots, A_{N} の中に、値 v が何個あるか

元々の数列  A_{1}, A_{2}, \dots, A_{N} の並びはどうでもよくて、各値が何個あるのかだけが重要なので、このような変換が有効になる。このような配列をバケットと呼ぶことがある。

バケットにして考えることが有効な類題として、次の問題がある (しかもこの類題は、後述する差分更新の考え方をする、という点でも似ている)

atcoder.jp

差分更新

改めて、各クエリをどう処理するのかを考えよう。実は、このように変換してあれば簡単で


  • num[ C ] += num[ B ] とする (このとき、総和は C × num[ B ] だけ増加する)
  • num[ B ] を 0 にする (このとき、総和は B × num[ B ] だけ減少する)

という風にするだけだ。そして、この前後で、数列の総和は、

  • (C - B) × num[ B ] だけ増加

することがわかる。以上から、最初に数列 A の総和を求めておいて、その後は上記のようなクエリ処理を行いながら、総和も更新していけば OK。

計算量は  O(N + Q) となる。

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

int main() {
    int N, Q;
    cin >> N;
    long long sum = 0;
    vector<long long> num(110000, 0);
    for (int i = 0; i < N; ++i) {
        long long A; cin >> A;
        sum += A;
        num[A]++;
    }
        
    cin >> Q;
    for (int i = 0; i < Q; ++i) {
        long long B, C; cin >> B >> C;
        sum += (C - B) * num[B];
        num[C] += num[B];
        num[B] = 0;
        cout << sum << endl;
     }
}