torus711 のアレ

主に競技プログラミングの問題について書きます.PC 以外だと数式が表示されないかもしれないです

AtCoder Beginner Contest 324, E : Joint Two Strings

問題概要

 英小文字からなる $N$ 個の文字列 $\{ S_i \}_{ 0 \leq i < N }$ と,同じく英小文字からなる文字列 $T$ が与えられる.
 $0 \leq i, j < N$ を満たす 2 つの整数 $i, j$ であって,$S_i + S_j$ が $T$ を部分列として含むようなものの個数はいくつか? ここで,$+$ は文字列の連結を表す.

制約

  • $N \in \mathbb Z^+$
  • $1 \leq N \leq 5 \times 10^5$
  • $\sum |S_i| \leq 5 \times 10^5$

解法

 すぐに思い付く愚直解法としては $i, j$ の組合せを全て試すというものがありますが,この制約だと $\Omega( N^2 )$ 時間のアルゴリズムでは TLE してしまうので,より効率的な方法を考える必要があります.
 ではどういうことができたらうれしいかというと,$i$ だけを固定したときに条件を満たす $j$ の個数をある程度高速に求められれば,問題を解くことができます*1.そのために,題意で要求されている条件を使いやすい形に言い換える必要があります.
 本解法では,$S_i$ が $T$ の手前側を,$S_j$ が $T$ の後ろ側を(それぞれ空かもしれませんが)カバーするという性質を利用します.$S_i$ が $T$ を先頭から $a_i$ 文字「消費」*2できて,$S_j$ が $T$ の後ろ $b_j$ 文字を消費できるとします.このとき,$a_i + b_j \geq |T|$ であれば,$T$ 全体を消費できます.
 消費できる文字数の求め方は後述しますが,$i$ を固定すると $a_i$ も固定されます*3.この条件下で,条件を満たす $b_j$ を数えたいわけですが,そのような $b_j$ は $b_j \geq |T| - a_i$ という条件を満たします.ある値以上の値の個数を求めめたいという話なので,全ての $j$ について $b_j$ を求めてソートしておけば,二分探索によって $O( \log N )$ 時間で条件を満たすものの個数を数えられます*4
 では,消費できる文字数を計算する方法を考えたい訳ですが,後ろからの消費数というのは両文字列を反転すれば先頭からの場合に帰着できるので,先頭から何文字消費できるかだけ考えます.方針としては $T$ の各文字を $S_i$ のどの文字に対応付けるかということを考える訳ですが,対応付け可能な文字(というか位置)が複数ある場合は,より手前のものを採用した方が後ろでの自由度が増えて得なので,そのようにします.なので手順としては,$T$ を先頭から舐めて,$S_i$ で最後に対応付けた文字(一度も対応付けてないなら先頭)より後ろの文字の内,最初に一致するものと対応付ければよいです.
 これで必要な情報は揃ったので,

  1. 各 $b_j$ を求めてソートしておく.
  2. 各 $a_i$ について,二分探索によって条件を満たす $b_j$ の個数を求める.

という手順で問題を解くことができます.
 計算量についてですが,各 $a_i, b_j$ を求めるときに $O( | S_i | )$ 時間がかかります.$T$ の文字数はいいのか,という話がありますが,対応付けられる文字が見つからなければ break するようにすれば,$O( | S_i | )$ 時間になります.これはトータルで $O( \sum_{ 0 \leq i < N } | S_i | )$ 時間です.更に,1. のソート部分で $O( N \log N )$ 時間がかかり,2. で二分探索を $N$ 回行うので同様に $O( N \log N )$ 時間がかかります.なので全体では $O( \sum_{ 0 \leq i < N } | S_i | + N \log N )$ 時間となり,TL に間に合います.

コード

int main()
{
	cin.tie( nullptr );
	ios::sync_with_stdio( false );
	cout << setprecision( 12 ) << fixed;

	IN( int, N );
	IN( string, T );
	VS S( N );
	cin >> S;

	const int L = SZ( T );

	VI forward( N );
	REP( i, N )
	{
		int cnt = 0, pos = 0;
		FOR( c, T )
		{
			const auto p = S[i].find( c, pos );
			if ( p != string::npos )
			{
				++cnt;
				pos = p + 1;
			}
			else
			{
				break;
			}
		}
		forward[i] = cnt;
	}

	reverse( ALL( T ) );
	VI backward( N );
	REP( i, N )
	{
		reverse( ALL( S[i] ) );
		
		int cnt = 0, pos = 0;
		FOR( c, T )
		{
			const auto p = S[i].find( c, pos );
			if ( p != string::npos )
			{
				++cnt;
				pos = p + 1;
			}
			else
			{
				break;
			}
		}
		backward[i] = cnt;
	}

	sort( ALL( backward ) );

	LL res = 0;
	FOR( a, forward )
	{
		const int b = end( backward ) - lower_bound( ALL( backward ), L - a );
		res += b;
	}

	cout << res << endl;

	return 0;
}

*1:「片側だけを固定する」アイデアはどこから来たのかという話もありますが,このパターンは割とよく見かけるので,そういうものとしか…….

*2:もう少しちゃんと言うと,$S_i$ が $T$ の先頭 $a_i$ 文字を部分列として含むということ.

*3:最大でないような消費文字数をとっても損しかしないので考えません

*4:より正確には,条件を満たすようになる境界位置を求められるという話なので,そこから個数を求めるには適切なデータ構造(e.g. std::vector)で保持しておく必要があります.