AtCoder Beginner Contest 183, F : Confluence

問題文

　https://atcoder.jp/contests/abc183/tasks/abc183_f

問題概要

　$N$ 人の人がいて，人 $i$ はクラス $C_i$ に属している．以下の 2 種のクエリを処理せよ．

• クエリ 1 : 人 $a, b$ が含まれる集団が合流する
• クエリ 2 : 人 $x$ が属している集団で，クラス $y$ に属している人の人数を出力する

制約

• $1 \leq N, Q \leq w \times 10^5$
• その他題意と整合性のあるいい感じの制約

解法

　集団の併合だけならば，Union-Find で高速に処理することができます．なので，問題のキーポイントはクラスの情報をいかに管理するかです．
　ここで，クラスの情報を std::map< int, int > を使って，クラス→人数 な対応表を作って集団ごとにもっておくことを考えます．これを何も考えずに更新していくと TLE してしまいますが，Weighted -Union Heuristic を使うことで，各要素が移動する回数が $O( Q \log N )$ で抑えられ，オーダーレベルで高速化されます．

コード

class DisjointSetForest;
// DisjointSetForest( N )
// find( x )
// same( x, y )
// unite( x, y )
// groups()
// groupSize( x )

void merge( map< int, int > &a, map< int, int > &b )
{
	if ( SZ( a ) < SZ( b ) )
	{
		swap( a, b );
	}
	FOR( p, b )
	{
		a[ p.fst ] += p.snd;
	}
	b.clear();
	return;
}

int main()
{
	IN( int, N, Q );
	VI C( N );
	cin >> C;
	transform( ALL( C ), begin( C ), bind( minus< int >(), _1, 1 ) );

	VT< map< int, int > > maps( N );
	REP( i, N )
	{
		++maps[i][ C[i] ];
	}

	DisjointSetForest dsf( N );

	REP( Q )
	{
		IN( int, t, a, b );
		--a, --b;

		if ( t == 1 )
		{
			const int i = dsf.find( a );
			const int j = dsf.find( b );
			if ( i == j )
			{
				continue;
			}

			dsf.unite( a, b );
			merge( maps[i], maps[j] );

			const int k = dsf.find( a );
			if ( maps[k].empty() )
			{
				swap( maps[i], maps[j] );
			}
			assert( !maps[k].empty() );
		}
		else
		{
			auto &m = maps[ dsf.find( a ) ];
			cout << ( EXISTS( m, b ) ? m[b] : 0 ) << '\n';
		}
	}
	cout << endl;

	return 0;
}