AtCoder Beginner Contest 171, D : Replacing

問題文

　https://atcoder.jp/contests/abc171/tasks/abc171_d

問題概要

　$N$ 項からなる正整数列 $\{ A_i \}$ がある．ここから，次の操作を $Q$ 回行う．

• $i$ 回目の操作では，値が $B_i$ である要素をすべて $C_i$ に置き換える

　それぞれの $i$ について，操作後の $\sum A$ を求めよ．

制約

• $1 \leq N, Q, A_i, B_i, C_i \leq 10^5$
• $B_i \neq C_i$

解法

　値の置き換えをする度に配列の要素を舐めたりすると TLE します．しかし，興味があるのは和だけであり，関わりのある整数の種類は高々 $2$ 種であることから，全体の和と値ごとの個数を覚えておけば和の差分を求めることができます．
　個数を std::map などで管理することにすると，操作ごとに map を定数回操作するので，全体で $O( Q \log N )$ 時間となります．

コード

int main()
{
	IN( int, N );
	VI A( N );
	cin >> A;
	
	LL sum = accumulate( ALL( A ), 0LL );

	map< int, int > counts;
	for_each( ALL( A ), [&]( const int a ){ ++counts[a]; } );

	IN( int, Q );
	REP( Q )
	{
		IN( LL, B, C );

		sum -= B * counts[B];
		sum -= C * counts[C];
		counts[C] += counts[B];
		counts[B] = 0;
		sum += C * counts[C];

		cout << sum << '\n';
	}
	cout << flush;

	return 0;
}