ABC171のD問題に挑戦してみた - インフラエンジニアの備忘録とか色々

初めに

今回はABC171のD問題に挑戦してみました。
この問題は、解説を見ないで回答することができましたが、
解いていく中で気付きがあったのでそれを中心に記載して行きます。

※ブログ内で問題の制約事項をすべて記載はしないので、制約事項は↑のURLより参照願います。

D - Replacing

N個の整数 $A_1～A_N$ からなら数列が与えられます。その数列に対して1 → 2といった
数字の置き換え操作をQ回行い、1～Q回目の入れ替えのタイミングでの
数列の和をそれぞれ出力する問題です。

最初の回答

#include<bits/stdc++.h>

using namespace std;
using ll = long long;
const long long INF = 1LL<<60;

typedef pair<ll,ll> Pair;
int main(){

ll N,Q;

cin >> N;

vector<ll> A(N);

for(ll i = 0;i < N;++i) cin >> A[i];

cin >> Q;

ll B,C;
vector<ll> S(Q,0);

for(ll i = 0;i < Q;++i){

    cin >> B >> C;

    for(ll j = 0;j < N;++j){

        if(A[j] == B){
            A[j] = C;
        }

        S[i] += A[j];

    }

}

for(ll i = 0;i < Q;++i){

    cout << S[i] << endl;

}

return 0;

}

最終的な回答

#include<bits/stdc++.h>

using namespace std;
using ll = long long;
const long long INF = 1LL<<60;

typedef pair<ll,ll> Pair;
int main(){

ll N,Q;
ll sum = 0;

cin >> N;

ll A;
vector<ll> cnt(100001,0);

for(ll i = 0;i < N;++i){
    cin >> A;
    // 初期入力された数列の和を計算し格納
    sum += A;
    // バケットに入力された整数の個数を格納
    ++cnt[A];
}

cin >> Q;

ll B,C;
vector<ll> S(Q,0);

for(ll i = 0;i < Q;++i){

    cin >> B >> C;

    // 入れ替え操作が発生した際の差分(C-B)と入れ替わった数字の個数の積をsumに加算
    sum += (C - B)*cnt[B];
    // 入れ替え操作により入れ替えもとの数字の個数を入れ替え先の数字の個数に加算
    cnt[C] += cnt[B];
    // 入れ替え元の数字の個数を0に設定
    cnt[B] = 0;
    
    S[i] = sum;

}

for(ll i = 0;i < Q;++i) cout << S[i] << endl;

return 0;

}

まずこの問題は愚直に入れ替え操作を行うと $O(NQ)$ となり、
N,Qの入力の制約が10⁵のため、TLEになります。(最初の回答)