AtCoder Beginner Contest 172, C : Tsundoku

問題文

　https://atcoder.jp/contests/abc172/tasks/abc172_c

問題概要

　本のスタックが 2 つあって，片方には $N$ 冊の本が，他方には $M $ 冊の本が積まれている．$N$ 札の本の内，上から $i$ 番目は読むのに $A_i$ 分かかり，$M $ 冊の本の内上から $j$ 番目は $B_j$ 分かかる．
　いずれかのスタックの一番上の本を読んでからそれを取り除く操作を繰り返すことを考える．$K$ 分以内に読み切れる本の最大値はいくらか？

制約

• $1 \leq N, M \leq 2 \times 10^5$
• $1 \leq K \leq 10^9$
• $1 \leq A_i, B_i \leq 10^9$

解法

　「各スタックの一番上の本について所要時間が短い方を読む」という Greedy アルゴリズムをやると WA します（やらかしました）．これでは，重い本の下に軽い本が沢山あるようなケースをカバーできません．よって，もう少し面倒なことをやる必要があります．
　片方の山から読む本の冊数を決めたとき，もう片方の山からは，可能な限り多くの本を読むのが最適な戦略となります．従って，片方の山から読む冊数を決めたときに他方から読める冊数をある程度高速に求められれば，問題を解くことができます．ここで例えば，山 $B$ の上から $j$ 冊を読むときの所要時間が全て分かっているとすれば，この値は明らかに単調増加なので，二分探索によって可能な最大値を求められます．また，この値は $B$ の累積和として求まる列であるので，事前に計算しておくことができます．
　この方針の場合，まず累積和の計算に $\Theta( M )$ 時間がかかります．その後，山 $A$ から読む冊数の候補が $\Theta( N )$ 通りあって，それぞれについて二分探索をして $O( \log M )$ 時間をかけるので，全体で $O( M + N \log M )$ 時間となります．

コード

int main()
{
	IN( int, N, M, K );
	VT< LL > A( N ), B( M );
	cin >> A >> B;

	VT< LL > psum_a( 1 ), psum_b( 1 );
	partial_sum( ALL( A ), BI( psum_a ) );
	partial_sum( ALL( B ), BI( psum_b ) );

	LL res = 0;
	REP( i, N + 1 )
	{
		const int j = upper_bound( ALL( psum_b ), K - psum_a[i] ) - begin( psum_b );
		if ( j )
		{
			chmax< LL >( res, i + j - 1 );
		}
	}
	cout << res << endl;

	return 0;
}