AtCoder Beginner Contest 340, E : Mancala 2

問題文

　https://atcoder.jp/contests/abc340/tasks/abc340_e

問題概要

　$0$ から $n - 1$ で番号付けられた $n$ 個の箱があり，最初，箱 $i$ には $A_i$ 個のボールが入っている．
　$i = 1, 2, \dots, m $ に対し順に以下の手続きを実行する．

1. 変数 $c$ を用意し，$c \leftarrow 0$ とする．
2. 箱 $B_i$ の中のボールを全て取り出し，手に持つ．
3. 手にボールをひとつ以上持っている間，次の処理を繰り返す．
   1. $c \leftarrow c + 1$ とする．
   2. 手に持っているボールをひとつ，箱 $( B_i + c ) \bmod n$ に入れる．

　操作完了後に各箱に入っているボールの個数を求めよ．

制約

• $1 \leq n, m \leq 2 \times 10^5$
• $0 \leq A_i \leq 10^9$
• $0 \leq B_i < n$

解法

　ボールの数がかなり多いので，提示された手続きを愚直に実行してボールをひとつずつ扱うようなことはできません．複数のボールに対する操作をまとめて行えるような工夫が必要になります．
　手続きについて観察してみましょう．困ったことが起こるのはボールの数が多いときなわけですが，例えば $n$ 個以上のボールを持っているとき，$( B_i + c ) \bmod n$ の値は一周して戻ってきます．その後もまだ $n$ 個以上持っていれば，同様のことが繰り返されます．これは，最初に手に持ったボールの個数を $b$ とすれば，

• $A$ の要素全てを $1$ 増やす操作

を $\left\lfloor \frac b n \right\rfloor$ 回繰り返すことだと言い換えられます．これは

• $A$ の要素全てを $\left\lfloor \frac b n \right\rfloor$ 増やす操作

を一回行っていると同じです．この操作の後，手の中のボールの個数は $b' = b \bmod n$ 個になっています．$b'$ もそれなりに大きな値をとり得るので，残ったボールについても言い換えを考えます．
　残ったボールを入れ始める位置は $$l = ( B_i + 1 ) \bmod n$$ で，入れ終わりの次の位置*1は $$r = ( l + b' ) \bmod n$$ です．ここで $l \leq r$ ならば，操作先は連続したひとつの区間 $$[ l, r )$$ になり，そうでないときは右にはみ出して左に戻るので，2 つの連続した区間
\begin{align*}
[ l, n ) \\
[ 0, r )
\end{align*}
になります．
　以上をまとめると，いずれの操作も

• $A$ 上の連続した区間に一様に値を加える

という形をしています．このような操作は遅延伝播セグメント木がサポートしているので，これを用いれば問題を解くことができます．実装については自分で用意してもよいですし，既にリンクしているように AtCoder Library (ACL) に含まれているので，これを利用することもできます．
　計算量についてですが，各 $i = 1, 2, \dots, m $ についてセグメント木を定数回ずつ操作するので，$O( m \log n )$ 時間がかかります．初期化や最後の集計は $O( n \log n )$ 時間なので，合わせると $O( ( n + m ) \log n )$ 時間となります．

コード

#include <atcoder/lazysegtree>
using namespace atcoder;

struct S
{
	LL v;
};

struct F
{
	LL a, b;
};

S op( S, S )
{
	return S{ -1 };
}

S e()
{
	return S{ 0 };
}

S mapping( F f, S s )
{
	return S{ f.a + s.v };
}

F composition( F f1, F f2 )
{
	return F{ f1.a + f2.a, -1 };
}

F id()
{
	return F{ 0, -1 };
}

int main()
{
	IN( int, N, M );
	VI A( N ), B( M );
	cin >> A >> B;

	lazy_segtree< S, op, e, F, mapping, composition, id > segtree( N );
	REP( i, N )
	{
		segtree.set( i, S{ A[i] } );
	}

	REP( i, M )
	{
		const LL balls = segtree.get( B[i] ).v;
		segtree.set( B[i], S{ 0 } );

		const LL whole = balls / N;
		const LL rest = balls % N;

		segtree.apply( 0, N, F{ whole, -1 } );

		const int L = ( B[i] + 1 ) % N;
		const int R = ( L + rest ) % N;
		if ( L <= R )
		{
			segtree.apply( L, R, F{ 1, -1 } );
		}
		else
		{
			segtree.apply( L, N, F{ 1, -1 } );
			segtree.apply( 0, R, F{ 1, -1 } );
		}
	}

	VLL res;
	REP( i, N )
	{
		res.PB( segtree.get( i ).v );
	}
	cout << res << endl;

	return 0;
}

*1:区間を半開区間で扱うので次の位置を採用します．