AtCoder Beginner Contest 258, E : Packing Potatoes

問題文

　https://atcoder.jp/contests/abc258/tasks/abc258_e

問題概要

　$N$ 種類の荷物があり，種類 $i$ の重さは $W_i$ である．
　今，荷物が無限に運ばれてきて，$i$ 番目 (0-indexed) の荷物の種類は $i \bmod N$ である．これらの荷物を次のように処理する．

1. はじめ，空の箱を一つ用意する．
2. 各荷物について，順に次の処理をする
   1. 当該荷物を箱に入れる
   2. その後，箱の内容物の重量が $X$ 以上になっていればその箱を除外し，新たな空箱を用意する

　次の形式で表されるクエリが $Q$ 個与えられるので，処理せよ：

1. 整数 $K$ が与えられる
2. $K$ 番目 (0-indexed) に除外された箱に詰められた荷物の個数を出力する

制約

• $1 \leq N, Q \leq 2 \times 10^5$
• $1 \leq X \leq 10^9$
• $1 \leq W_i \leq 10^9$
• $1 \leq K \leq 10^{ 12 }$

解法

　色々と値が大きいので難しそうに見えます．$K$ が非常に大きいので箱のリストを実際にもつことはできませんし，$N, X$ が大きいので箱に荷物を一つずつ詰めることも全ての箱を保持することもできません．それぞれ工夫する必要があります．
　記法についての注意ですが，クエリ処理のような文脈で，前処理に $O( f( * ) )$ 時間，クエリ処理に $O( g( * ) )$ 時間かかることを $\langle O( f( * ) ), O( g( * ) ) \rangle$ などと書くことにします．

箱について

　ちょっとよく考えると分かることとして，箱は「どの種類の荷物から詰め始めたか」によって特徴づけられるので，本質的には $\Theta( N )$ 種類しかありません．各 $i$ ($0 \leq i < N$) を開始位置として箱詰めしたときに作られる箱に関する情報を高速に計算できれば役立ちそうです．
　位置 $i$ から始めて $n$ 個の荷物を（$X$ に関する条件を無視して）選択する場合を考えてみます．このとき，その重みの総和は $n$ に関して単調性をもつので，二分法を用いれば $n$ の上界に対して対数オーター通りを試せば最適な $n$ を求められます．では $n$ を固定したとして，そのときの重みの総和はどうなるでしょうか？　$N$ 回のピックをまとめて一回の操作だと思うことにすれば，重みに対する寄与は，

• $N$ 個の荷物をまとめてピックする
• 位置 $i$ から始めて列 $W$ の末尾に到達するまでの重み
• $W$ の末尾を通り越して先頭に戻ってきてからいくつかの荷物の重み

の 3 通りに分類できます（それぞれ，1 つ以上あるとは限らない）．$N$ 個の要素をまとめてピックできる回数は $\lfloor \frac n N \rfloor$ となり，それで拾いきれない荷物の個数を $n'$ とすれば $n' = n \bmod N$ です．これらを使って上記の箇条書きをちゃんと書くと，

• $\sum W \times \lfloor \frac n N \rfloor$
• $\sum_{ j = i }^{ \min( N, i + n' ) - 1 } W_j$
• $\sum_{ j = 0 }^{ n' - ( N - i ) - 1 } W_j$ ($0 < n' - ( N - i )$ のとき)

となります．$W$ 上の区間の和が現れていますが，これは累積和を使うと $\langle \Theta( N ), O( 1 ) \rangle$ になるので，各 $n$ に対してそれぞれ $O( 1 )$ になって二分法が間に合います（上界が $X$ であることに注意する）．これで求まる情報に $A$ という名前をつけ，$$A_i = \text{位置 $i$ から箱詰めを始めるときに詰められる荷物の個数}$$ とします．
　また，添字に関して加算と剰余算で「次の箱に最初に入る荷物の添字」が求まります．これはあとで使うので，列 $B$ と呼ぶことにします．

$K$ 番目の箱を求める

　前節で求めた $B$ によって，「位置 $i$ から箱一つ分の箱詰めをした次の箱の開始位置」が求まります．$B$ に対する「累乗」のような概念を，次のように定義します：

• $\mathcal B^0_i = i$
• $\mathcal B^k_i = B_{ \mathcal B^{ k - 1 }_i }$ ($k > 0$)

コンセプトというか意味としては，$$\mathcal B^k_i = \text{位置 $i$ から始めて $k$ 個の箱に詰めたときの，次の箱の開始位置}$$ です．このような累乗（というかある種の関数合成の冪乗）はダブリングと呼ばれる手法によって高速に計算できます．今回の場合，$\langle \Theta( N \log { K_{ \mathit{ max } } } ), O( \log { K_{ \mathit{ max } } } ) \rangle$ 時間です．

答えの計算

　以上のことから，$K$ が与えられたときの $K$ 番目の箱に最初に詰める荷物を指す位置と，その位置から詰め始めたときに箱に詰める荷物の個数が求まります．ちゃんと書けば，$$A_{ \mathcal B^K_0 }$$ です．

計算量

　計算量については，面倒くさいので上界にだけ着目することにしつつ上に書き散らしているのをかき集めれば，

• 累積和の前計算に $O( N )$ 時間
• $N$ 回の二分法に $O( N \log X )$ 時間
• ダブリングの前計算に $O( N \log K_{ \mathit{ max } } )$ 時間
• $Q$ 回のクエリ処理でクエリごとに，
  • ダブリングのクエリ処理に $O( \log K_{ \mathit{ max } } )$ 時間

となり，全体では $O( N \log X + ( N + Q ) \log( K_{ \mathit{ max } } ) )$ 時間となります．

コード

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

	IN( int, N, Q, X );
	VT< LL > W( N );
	cin >> W;

	VT< LL > psum( 1 );
	partial_sum( ALL( W ), BI( psum ) );

	VI nexts( N );
	VT< LL > amounts( N );
	REP( i, N )
	{
		int lb = 0, ub = X;
		while ( lb + 1 < ub )
		{
			const int mid = ( lb + ub ) / 2;

			LL w = 0;
			w += psum[N] * ( mid / N );
			const LL rest = mid % N;
			w += psum[ min< LL >( N, i + rest ) ] - psum[i];
			if ( 0 < rest - ( N - i ) )
			{
				w += psum[ rest - ( N - i ) ];
			}
			( w < X ? lb : ub ) = mid;

		}
		nexts[i] = ( i + ub ) % N;
		amounts[i] = ub;
	}

	VVI dp( 40, VI( N ) );
	dp[0] = nexts;

	REP( i, 39 )
	{
		REP( j, N )
		{
			dp[ i + 1 ][j] = dp[i][ dp[i][j] ];
		}
	}

	REP( Q )
	{
		IN( LL, K );
		--K;

		int res = 0, k = 0;
		for ( ; K; K >>= 1, ++k )
		{
			if ( K & 1 )
			{
				res = dp[k][ res ];
			}
		}

		cout << amounts[ res ] << '\n';
	}

	cout << flush;

	return 0;
}