ABC340 E - Mancala 2

atcoder.jp

問題

• 円状に並んだN個の箱があるよ
• 箱iには、$A_{i}$個のボールが入ってるよ
• 1~Mまで、以下の操作を行うよ
  • 変数Cを0とする
  • 箱$B_i$の中のぼーるを全部出して、箱$B_i+1$から順に1つずつ入れていくよ
• すべての操作を終えた時に、各箱に入っているボールの数を教えて

成約

• $1 \leq N \leq 5 \times10^{5}$
• $0 \leq D \leq 5 \times10^{5}$
• $1 \leq A_{i} \leq 5 \times10^{5}$

思考

• 操作を丁寧に考える
• 箱の中のボールを全部出すのは、その箱の中にあるボールの数を0にする操作
• 箱$B_{i}+1$から順に1つずつ入れていく操作は…区間に一気に足す操作
  • 1個ずつシミュレーションしてたんじゃ間に合わないので…とりあえず、1週（N個）分は、まとめて捌いて…
  • 区間更新(ここからここまでの要素全てに+1）と1点取得ができれば幸せになれそう
    • 遅延セグ木？

コード

from atcoder.lazysegtree import LazySegTree

N,M = map(int,input().split())
A_L = list(map(int, input().split()))
B_L = list(map(int, input().split()))

LST = LazySegTree(lambda a,b:a+b,0,lambda a,b:a+b,lambda a,b:a+b,0,A_L)

for box in B_L:
    ball_cnt = LST.get(box)
    LST.set(box,0)
    if ball_cnt >= N:
        add_cnt = ball_cnt//N
        LST.apply(0,N,add_cnt)
        ball_cnt %= N

    if box + 1 + ball_cnt < N:
        LST.apply(box+1,box+1+ball_cnt,1)
    else:
        LST.apply(0,N,1)
        LST.apply((box+1+ball_cnt)%N,box+1,-1)

ans_L = []
for i in range(N):
    ans_L.append(LST.get(i))

print(*ans_L)