torus711 のアレ

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

TopCoder, Single Round Match 654, Division 1, Level 1 : SquareScores

問題概要

 文字列 S について,その文字列の得点を次のように定める

  • インデックスの対 i, j であって,部分文字列 S[ i, j ] がただ一種類の文字から成るようなものの個数

 今,文字列 S が与えれるが,S の内いくつかの文字は決まっていない('?' で表される).各 '?' を(独立に)アルファベットに置換する.置換後の文字は確率的に選択される.配列 p が与えられ,ある '?' を置換するとき,文字 i 番目のアルファベットに置換される確率は p_i パーセントである.
 全ての '?' を置換し終わった状態でのスコアの期待値を求めよ.
 |S| \leq 1,000


 期待値の線形性を使う解法と,動的計画法による解法の 2 通りについて書きます.
 また,どちらの記述に於いても,|S|L と書きます.

解法(期待値の線形性)

 求めたい期待値は,部分文字列 S[ i, j ] から得られる得点を e( i, j ) とすると E[ \sum_{ i = 0 }^{ L - 1 }\sum_{ j = i }^{ L - 1 } e( i, j ) ] です(得られる得点が 1 ずつなので確率だけの和です).期待値の線形性を使うとこの式を \sum_{ i = 0 }^{ L - 1 }\sum_{ j = i }^{ L - 1 } E[ e( i, j ) ] と変形できます.つまり,s の全ての部分文字列(異なる位置に出現するものは全て相異なるものと扱う)について,その部分文字列が条件を満たす確率を求めて足し合わせれば答えが求まります.
 さて,s の部分文字列が 1 つ与えられたとき,'?' を題意に従って置換した結果,部分文字列全体が同じ文字になる確率を考えていきます.まず,どの文字で一致するかによって違う事象となることを気を付けないといけませんが,異なる文字になる事象というのは互いに素な事象なので,全体の確率を求めるときには単純に加算してしまってよいです.そうすると,ある s の部分文字列と文字 c が与えられたとき,部分文字列が全て c になる確率を求められれば,いずれかの文字で埋まる確率を求められます.
 当然ですが,部分文字列中に c でない文字が存在する場合はどう置換しても全部 c にはならないので確率は 0 です.そうすると,確率が正になり得る場合というのは,部分文字列が c 及び '?' のみからなる場合です.この場合,'?' が全て c に変化すればよいので,全てが c になる確率は '?' の数を qck 番目のアルファベットであるとして, { p_k }^q となります.
 上記の確率を計算するために必要な情報は次の 2 つです.

  • 区間内に存在する '?' の個数
  • 区間内に存在する文字が一種類か否かの判別を可能にする情報

'?' の個数については普通に数えてから累積和をとることで,任意の区間に対し O( 1 ) で求められます.また,後者についても,各文字が区間中にいくつあるかを求められるようにしておくことで,(着目文字の個数)+('?' の数) と区間の長さを比較することで求められます(文字が一種類なら等しくなる).結局,各文字についても個数の累積和を前計算しておくことで,後者も O( 1 ) で求まります.あとは,区間のとり方及び文字の組合せ全てについて,上記のように確率を求めて足し合わせれば答えが求まります.
 累積和の前計算は O( L ) 時間で実行でき,答えを求める部分の処理時間は部分文字列の総数に比例するので O( L^2 ) 時間となります.

コード

using VI = vector< int >;
using VVI = vector< vector< int > >;
template < typename T = int > using VT = vector< T >;

#define REP2( i, n ) REP3( i, 0, n )
#define REP3( i, m, n ) for ( int i = ( int )( m ); i < ( int )( n ); ++i )
#define GET_REP( a, b, c, F, ... ) F
#define REP( ... ) GET_REP( __VA_ARGS__, REP3, REP2 )( __VA_ARGS__ )
#define ALL( c ) ( c ).begin(), ( c ).end()

#define SZ( v ) ( (int)( v ).size() )
#define BI back_inserter

class SquareScores
{
public:
	double calcexpectation( vector <int> p, string s )
	{
		const int L = SZ( s );

		VT< double > pr( 26 );
		transform( ALL( p ), pr.begin(), bind( divides< double >(), _1, 100 ) );

		VI qcsum( 1, 0 );
		{
			VI counts;
			transform( ALL( s ), BI( counts ), bind( equal_to< char >(), _1, '?' ) );
			partial_sum( ALL( counts ), BI( qcsum ) );
		}

		VVI csums( 26, VI( 1, 0 ) );
		{
			VVI counts( 26 );
			for ( char c = 'a'; c <= 'z'; ++ c )
			{
				transform( ALL( s ), BI( counts[ c - 'a' ] ), bind( equal_to< char >(), _1, c ) );
			}
			REP( i, 26 )
			{
				partial_sum( ALL( counts[i] ), BI( csums[i] ) );
			}
		}

		double res = 0;
		REP( i, L )
		{
			REP( j, i, L )
			{
				REP( c, 26 )
				{
					const int len = j - i + 1;
					const int ques = qcsum[ j + 1 ] - qcsum[i];
					const int cnt = csums[c][ j + 1 ] - csums[c][i];

					if ( ques + cnt != len )
					{
						continue;
					}

					res += pow( pr[c], ques );
				}
			}
		}

		return res;
	}
};

解法(動的計画法

 まずスコアの計算について考えてみます.ある文字が n 文字連続している文字列があったとき,この文字列から得られるスコアは部分文字列の取り出し方の総数に等しくなります.n 文字の文字列から部分文字列を取り出す方法の総数は,\frac{ n( n + 1 ) }{ 2 } 通りです.これは,1, 2, 3, \dots, n 文字の部分文字列をそれぞれ n, n - 1, n - 2, \dots, 1 箇所から取れると考えると分かりやすいと思います.更に,この文字列にに更に同じ文字を付け加えて n + 1 文字になった状況を考えます.このときのスコアの増分は,1, 2, 3, \dots, n + 1 文字の部分文字列を取れる位置が 1 ずつ増える(全て末尾を左派時に含む位置)ので,n + 1 点増加します.あるいは真面目に計算して,\frac{ ( n + 1 )( n + 2 ) }{ 2 } - \frac{ n( n + 1 ) }{ 2 } = \frac{ n^2 + 3n + 2 - n^2 - n }{ 2 } = \frac{ 2n + 2 }{ 2 } = n + 1 です.この計算に必要なのは,文字を付け加えた後の状態に於いて,末尾の文字が何文字連続しているかという情報のみです.付け加える前の状態から見ると,付け加える文字が決まったとき,末尾の文字・連続している数の 2 つの情報があれば,付け加えた後で末尾の文字が何文字連続しているかを計算できます.この事を利用すると,1 文字ずつ文字列を伸ばしながらスコアを計算することができるようになります.また,'?' の扱いは,各状態に於いてそこに至る確率を持ちながら文字列を伸ばせば,遷移先の状態に至る確率を計算できます.確率が分かっていれば,スコアとの積をとることで期待値を計算できます.
 'しかし,単純に DFS をすると(アルファベットの種類 = 26 通りの選択肢が最大 L 回あるので)O( 26^L ) 時間かかってしまいます.そこで,次のような動的計画法を考えます.

  • dp[ 決めた文字数 ][ 末尾の文字(に対応する値) ][ 末尾の文字が連続している数 ] := 確率

初期状態では dp[ 0 ][ 26 ][ 0 ] のみ 1 で,それ以外は 0 です(26 はどのアルファベットにも対応しない値.先頭をいい感じに使うための一般的なテク).dp[ i ][ j ][ k ] からの遷移は,次に付け加える文字を全通り試すことになります.c 番目のアルファベットを付け加えるとすると,

  • dp[ i + 1 ][ c ][ nk ] を更新

となります.ここで,nk は j と c が等しいければ k + 1 で,それ以外のとき 1 です.また,遷移が発生する確率も 3 通りに場合分けされ,

  • p_c ( s[ i ] = '?' のとき)
  • 1 ( s[ i ] - 'a' = c のとき)
  • 0 ( それ以外 )

となり,現在の状態の確率に,この確率を掛けた値を遷移先に足し込めばよいです.
 DP が終わってから各 j, k に対して dp[ L ][ j ][ k ] * k の総和を求めれば答えが求まりますが,普通にやってしまうと空間計算量が定数大きめの O( L^2 ) となってメモリが若干厳しいです.そこで,一つ目のパラメータが 1 ずつしか増加しないことを利用して,i の奇偶で場合分けして表を使い回すアレをやります.こうすると答えを後から集計できないので,DP しながら答えを求めます(求め方は変わりません).
 DP の状態数が定数大きめの O( L^2 ) であり,状態遷移も大きめの定数なので,全体は定数大きめの O( L^2 ) 時間で実行できます.ただ,「定数大きめ」としつこく書いていることからお察しの通り,何も工夫しないと TLE します.そこで,確率が無視できるほど小さい状態からの処理をスキップすることで,状態遷移側の定数を小さくする工夫をします.10^{ -15 } 以下になったら無視するという処理をしたところ,最大 200 ms 程度で通りました.

コード

template < typename T = int > using VT = vector< T >;

#define REP2( i, n ) REP3( i, 0, n )
#define REP3( i, m, n ) for ( int i = ( int )( m ); i < ( int )( n ); ++i )
#define GET_REP( a, b, c, F, ... ) F
#define REP( ... ) GET_REP( __VA_ARGS__, REP3, REP2 )( __VA_ARGS__ )
#define ALL( c ) ( c ).begin(), ( c ).end()
#define AALL( a, t ) ( t* )a, ( t* )a + sizeof( a ) / sizeof( t )

#define SZ( v ) ( (int)( v ).size() )

double dp[ 2 ][ 32 ][ 1024 ];

class SquareScores
{
public:
	double calcexpectation( vector <int> p, string s )
	{
		{ // initialize for local test
			fill( AALL( dp, double ), 0 );
		}

		const int L = SZ( s );

		VT< double > pr( 26 );
		transform( ALL( p ), pr.begin(), bind( divides< double >(), _1, 100 ) );

		double res = 0;
		dp[0][ 26 ][0] = 1;

		REP( i, L )
		{
			const int cur = i % 2;
			const int next = 1 - cur;

			fill( AALL( dp[ next ], double ), 0 );

			REP( j, 27 )
			{
				REP( k, 1001 )
				{
					if ( dp[ cur ][j][k] < 1e-15 )
					{
						continue;
					}

					REP( c, 26 )
					{
						const int nk = j == c ? k + 1 : 1;
						const double q = s[i] == '?' ? pr[c] : ( s[i] - 'a' == c );
						dp[ next ][c][ nk ] += dp[ cur ][j][k] * q;
					}
				}
			}

			REP( j, 26 )
			{
				REP( k, 1001 )
				{
					res += dp[ next ][j][k] * k;
				}
			}
		}

		return res;
	}
};