AtCoder Regular Contest 104, B : DNA Sequence

問題文

　https://atcoder.jp/contests/arc104/tasks/arc104_b

問題概要

　`A', 'T', 'C', 'G' の 4 種の文字からなる文字列について考える．長さの等しい文字列 $T_1, T_2$ が「相補的」であるとは，各 $i$ について $\{ T_{ 1, i }, T_{ 2, i } \} = \{ \text{'A'}, \text{'T'} \}$ または $\{ T_{ 1, i }, T_{ 2, i } \} = \{ \text{'C'}, \text{'G'} \}$ であることを言う．
　長さ $N$ の文字列 $S$ が与えられる．$S$ の連続する*1部分文字列 $T$ であって，次の条件を満たすものの個数を求めよ．

• $T$ と相補的であるような，$T$ を並べ変えた文字列が存在する

ここで部分文字列は，取り出す位置が異なる場合は異なるものとする．

制約

• $1 \leq N \leq 5{,}000$
• $S_i \in \{ \text{'A'}, \text{'T'}, \text{'C'}, \text{'G'} \}$

解法

　まず，部分文字列 $T$ がどのような性質をもつべきかを考えます．$T$ が自身の並べ替えと相補的であるためには，各 $i$ について，$T_i$ と対応する文字（e.g. 'A' なら 'T', 'C' なら 'G'）を $T$ のどこかからもってこれる必要があります．ある一種の文字全てについて同じことが言えるので，例えば 'A' に着目すると 'T' の個数は 'A' の個数以上でなければなりません．同様に 'T' に着目すると 'A' の個数は 'T' の個数以上でなければならないので，まとめると，'A' の個数と 'T' の個数は等しくなければなりません．これと同じことが 'C', 'G' についても言えるので，結局，

• 'A' の個数と 'T' の個数が等しい
• 'C' の個数と 'G' の個数が等しい

ような部分文字列 $T$ を数えればよいことになります．
　$S$ の連続する部分文字列は，始点・終点のとり方の総数程度の個数で，$\Theta( N ^ 2 )$ 個です．このそれぞれについて，各文字の個数を（$\Omega( N^3 )$ 時間かけると TLE するので，）ある程度高速に求められれば，後は個数の比較だけで問題を解くことができます．
　始点の位置 $i$ が共通する部分文字列たちについて考えてみます．$i$ 文字目から $j - 1$ 文字目までの部分文字列を $S[ i, j )$ と書くことにします．まず，空文字列 $S[ i, i )$ については各文字の個数は明らかに $0$ なので $O( 1 )$ 時間で求まります．ある $j$ について $S[ i, j )$ の各文字の個数が分かっているとします．このとき，$S[ i, j + 1 )$ での各文字の個数は，$S_j$ の一文字分が加わるだけなので，$O( 1 )$ 時間で求められます．このように，始点を固定して右に伸ばしながら差分を求めていけば，各部分文字列あたり $O( 1 )$ 時間で文字数のカウントを更新できます．
　よって，アルゴリズムの全体は $\Theta( N ^ 2 )$ 時間となり，間に合います．

コード

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

	LL res = 0;
	REP( i, N )
	{
		VI counts( 256 );
		REP( j, i, N )
		{
			++counts[ S[j] ];

			const int a = counts[ 'A' ];
			const int t = counts[ 'T' ];
			const int c = counts[ 'C' ];
			const int g = counts[ 'G' ];

			res += a == t && c == g;
		}
	}
	cout << res << endl;

	return 0;
}

*1:本番で見落とした人です