AtCoder Beginner Contest 346, D : Gomamayo Sequence

問題文

　https://atcoder.jp/contests/abc346/tasks/abc346_d

問題概要

　$n$ 項からなるバイナリ列 $S = \langle S_1, S_2, \dots, S_n \rangle$ ($S_i \in \{ 0, 1 \}$) が与えられる．$S$ に対し，次の操作を任意回適用できる．

• コスト $C_i$ かけて $S_i$ を反転する．

　$S$ の添字 $i$ ($1 \leq i < n$) であって，$S_i = S_{ i + 1 }$ となるものが丁度ひとつであるようにするためにかかるコストの最小値はいくらか？

制約

• $2 \leq n \leq 2 \times 10^5$
• $1 \leq C_i \leq 10^9$

解法

　解を与えるバイナリ列を $\hat S$ と書くことにします．$\hat S$ がどのような形をしているかを考えると，$\hat S_i = \hat S_{ i + 1 }$ な一箇所以外では $0, 1$ が交互に出現するということなので，ちょっと言い換えると

• 前半は $01010\dots$ に一致し，後半は $10101\dots$ に一致
• 前半は $10101\dots$ に一致し，後半は $01010\dots$ に一致

のいずれかです．この前半・後半の境目は $\Theta( n )$ 通りなので，そのそれぞれについてある程度高速にコストを計算できれば，問題を解くことができます．
　$S$ の $i$ 文字目から $j$ 文字目までを $01010\dots$ に一致させるコストは，列 $\mathit{ cost }_{ 01 }$ を
\begin{equation*}
\mathit{ cost }_{ 01, i } = \begin{cases}
0 & \text{($S_i = 01010\dots_i$)} \\
C_i & \text{(otherwise)}
\end{cases}
\end{equation*}
とすれば，$\sum\limits_{ k = i }^{ j } \mathit{ cost }_{ 01, k }$ です．$10101\dots$ に一致させることについても，同様に $\mathit{ cost }_{ 10 }$ を定義すれば同じようにできます．列の部分和を求めるのは，「累積和」と呼ばれるテクを使えば，前計算に $\Theta( n )$ 時間をかけることで，任意の区間の和を $O( 1 )$ 時間で求められます．
　以上をまとめると，全体で $\Theta( n )$ 時間で問題を解けます．

コード

constexpr auto INF = LIM< LL >::max() / 2;

int main()
{
	IN( int, N );
	IN( string, S );
	VI C( N );
	cin >> C;

	VLL cost01( N ), cost10( N );
	REP( i, N )
	{
		cost01[i] = S[i] == ( '0' + i % 2 ) ? 0 : C[i];
		cost10[i] = S[i] == ( '1' - ( i % 2 ) ) ? 0 : C[i];
	}

	VLL psum1( 1 ), psum2( 1 );
	partial_sum( ALL( cost01 ), BI( psum1 ) );
	partial_sum( ALL( cost10 ), BI( psum2 ) );

	LL res = INF;
	REP( i, 1, N )
	{
		chmin( res, psum1[i] + ( psum2[N] - psum2[i] ) );
		chmin( res, psum2[i] + ( psum1[N] - psum1[i] ) );
	}
	cout << res << endl;

	return 0;
}