torus711 のアレ

主に競技プログラミングの問題について書きます

Codeforces #188, Division 2, B : Strings of Power

概要

英小文字からなる長さが 10^6 以下の文字列が与えられる。
この文字列の部分文字列の内、"heavy" をプレフィックスにもち "metal" をサフィックスにもつ部分文字列の数を求めよ。
(元の文字列に於ける出現位置が違うならば違う部分文字列であるとする)

解法

部分文字列の先頭と末尾を全通り試すと TLE してしまいます。
しかし、全通り試さなくても数え上げることができます。

有効な部分文字列の先頭とする "heavy" の位置を決めたとき、使うことができる有効な末尾位置は先頭位置以降にある "metal" の末尾です。
"heavy" と "metal" の出現位置を予め列挙しておくと、ある "heavy" に対して使うことのできる最も先頭に近い "metal" の位置は二分探索により高速に求めることができます。
出現位置は予め列挙されているので、最も早いものの位置が分かれば、それ以降のものの数は簡単に求めることができます。
一つの "heavy" に対する "metal" の数を O(N) で求められるので、全体では O(N \log N) で答えを求めることができます。
(しゃくとり法っぽくもできそう。その場合 O(N)

コード

typedef long long LL;
typedef vector<int> VI;

#define REP( i, m, n ) for ( int i = (int)( m ); i < (int)( n ); ++i )
#define ALL( c ) (c).begin(), (c).end()
#define PB( n ) push_back( n )

int main()
{
	cin.tie( 0 );
	ios::sync_with_stdio( false );

	string str;
	cin >> str;
	
	VI starts, ends;
	{
		int last = -1, pos;
		while ( ( pos = str.find( "heavy", last + 1 ) ) != string::npos )
		{
			starts.PB( pos );
			last = pos;
		}
		last = -1;
		while ( ( pos = str.find( "metal", last + 1) ) != string::npos )
		{
			ends.PB( pos );
			last = pos;
		}
	}

	LL res = 0;
	REP( i, 0, starts.size() )
	{
		auto pos = upper_bound( ALL( ends ), starts[i] );
		res += ends.end() - pos;
	}

	cout << res << endl;

	return 0;
}