ABC270E - Apple Baskets on Circle

周回回数を効率的に求め，残り$N$個以下になったら愚直に数えることを考える．

解法1：二分探索

考え方

Editorial - TOYOTA MOTOR CORPORATION Programming Contest 2022(AtCoder Beginner Contest 270)

「$K$未満になる最大のループ回数」を知れば，あとは1周以内で$K$個を食べられる．

「$x$回ループしたときに，食べるりんごは$K$個未満か？」という問題を二分探索で求める．

回答例

N, K = map(int, input().split())
A = list(map(int, input().split()))

def f(x):
    return sum(min(a, x) for a in A) < K

ok, ng = 0, max(A)
while ng - ok > 1:
    mid = (ok + ng) // 2
    if f(mid):
        ok = mid
    else:
        ng = mid

K -= sum(min(a, ok) for a in A) 
A = [max(0, a - ok) for a in A]
for i in range(N):
    if K == 0: break
    if A[i] > 0:
        A[i] -= 1
        K -= 1

print(*A)

解法2：シミュレーション

考え方

Editorial - TOYOTA MOTOR CORPORATION Programming Contest 2022(AtCoder Beginner Contest 270)

「$A$の一番小さい要素を使い切れるか？」に着目すると，$A$の要素数（$O(N)$）で周回回数を計算できる．

常に$A$の最小の要素を取り出すために，「優先度付きキュー」を使う．

• 「最もりんごが少ない$A$の要素」の回数だけループしたときに，$K$以上のりんごが取れるか判定する．
  
  • Trueの場合：ループ回数に取り出した要素数を加算，$K$に食べたりんごの数を加算する
  • Falseの場合：$\lfloor K / \text{(残りの要素数)}\rfloor$回だけループ回数を加算，$K$に食べたりんごの個数を加算する
  この処理を繰り返すと，$K < \text{(残りの要素数)} \leq N$となる．
  
• りんごが残っているかごを順に見ていき，合計$K$個食べるまで愚直に数える

回答例

from heapq import heapify, heappop

N, K = map(int, input().split())
A = list(map(int, input().split()))

A2 = [a for a in A]
heapify(A2)

loop = 0
rem = N
while (A2[0] - loop) * rem <= K:
    a = heappop(A2)
    K -= (a - loop) * rem
    loop += a - loop
    rem -= 1
    if K == 0: break
    if rem == 0: break

if rem > 0:
    add = K // rem
    K -= add * rem
    loop += add

A = [max(0, a - loop) for a in A]
for i in range(N):
    if K == 0: break
    if A[i] > 0:
        A[i] -= 1
        K -= 1

print(*A)