AtCoder Beginner Contest 345, C : One Time Swap

問題文

　https://atcoder.jp/contests/abc345/tasks/abc345_c

問題概要

　文字列 $S$ が与えられる．$S$ に対し，$S$ の添字 $i, j$ ($1 \leq i < j \leq |S|$) を選び，$S_i$ と $S_j$ を入れ替える操作をちょうど $1$ 回行う．この操作によって作られる文字列は何種類あるか？

制約

• $1 \leq |S| \leq 10^6$
• $S_i \in \mathcal C$

　ここで，$\mathcal C$ はすべての英小文字からなる集合とする．

解法

　操作により作られる文字列を $S'$ とし，文字 $c \in \mathcal C$ の $S$ 中の出現回数 ($| \{ i \in \{ 1, 2, \dots, |S| \} \mid S_i = c \} |$) を $\mathrm{ cnt }( c )$ とします．
　$S$ と $S'$ の関係にはふたつの場合があって，それは

• $S' = S$
• $S' \neq S$

です．前者は $\mathrm{ cnt }( c ) > 1$ なる $c \in \mathcal C$ が存在すれば可能です．後者はもう少し考える必要があって，異なる文字（を指す添字）を選ぶ方法の総数となります．文字 $c, d \in \mathcal C$ をとって，（辞書式順序の意味で）$c > d$ とすれば，文字 $c, d$ を交換して $S' \neq S$ なる $S'$ を生成する方法は $$\mathrm{ cnt }( c ) \mathrm{ cnt }( d )$$ 通りです．すべての valid な $c, d$ について和をとればよいので，$S' \neq S$ な $S'$ は $$\sum_{ c \in \mathcal C } \sum_{ d \in \mathcal C, c > d } \mathrm{ cnt }( c ) \mathrm{ cnt }( d )$$ 通りとなります．
　$\mathrm{ cnt }$ は，英小文字の ASCII コードが小さいことを利用すれば配列や std::vector で記憶できるので，$\Theta( |S| )$ 時間で計算できます．$\mathrm{ cnt }( c ) > 1$ な $c$ は $\Theta( | \mathcal C | )$ 時間で探索でき，$c, d \in \mathcal C$ の取り方をすべて試す処理は $\Theta( | \mathcal C |^2 )$ 時間で実行できます．「$\mathcal C$ は英小文字」ということから $| \mathcal C | = 26$ は定数だという立場をとるならば，この解法は $\Theta( |S| )$ 時間で走ります．

コード

int main()
{
	IN( string, S );

	VLL counts( 256 );
	for_each( ALL( S ), [&]( const char c ){ ++counts[c]; } );

	LL res = 0;
	bool same = false;
	for ( char c = 'a'; c <= 'z'; ++c )
	{
		same |= 2 <= counts[c];
		for ( char d = 'a'; d < c; ++d )
		{
			res += counts[c] * counts[d];
		}
	}

	res += same;

	cout << res << endl;

	return 0;
}