はじめに
HHKBプログラミングコンテスト 2025 (AtCoder Beginner Contest 388) にトライしました. A, B, C の 3 問は解けました. しかし D 問題はリストの特定範囲に定数を加算する効率的な方法が分からず, TLE (実行制限時間超過) で終了しました.
解説には いもす法 を使うと書かれてたので, 復習しました.
いもす法
いもす法 を読んで理解してください.
ABC388-D の考え方
コーディング時のリストと合わせるため, 宇宙人のナンバリングは $0$ から $N-1$ とします. また宇宙人 $i$ が持っている石の数を A[i] とします.
渡す石の数
もし宇宙人 $i$ が十分な石を持っていれば, 宇宙人 $i+1$ から $N-1$ までの $N-1-i$ 人に石を渡します.
しかし宇宙人 $i$ が十分な石を持っていなければ, 自分の持っている石の数が上限になります.
そのため, 渡す石の数は以下のように表現できます.
stone_give = min(A[i], N-1-i)
もらう石の数
宇宙人 $i$ がもらう石の数を C[i], その差分を D[i] = C[i] - C[i-1] とします.
宇宙人 $i$ が成人すると, i+1 から i+1+stone_give の宇宙人に石を渡します. そのため, いもす法 の考え方で D[i+1] を 1 増やし, D[i+1+stone_give] を 1 減らします. (下のコードでは, 念のため, 範囲が $N$ を越えないように D[min(N, i + stone_give + 1)] を 1 減らしています)
D[i+1] += 1 # D[] を 1 増やす
D[min(N, i + stone_give + 1)] -= 1 # stone_give+1 だけ先の D[] を 1 減らす
D を更新できたら, 累積和の考え方で C を更新し, 宇宙人 $i$ の石の数 A[i] を更新します.
if i != 0: # i=0 のときは石はもらわない
C[i] = C[i-1] + D[i]
A[i] += C[i]
コード
以上を用いてコーディングします.
ABC388_D.py
N = int(input())
A = list(map(int, input().split()))
C = [0] * N # C[i] ; A[i] が受け取る石の数
D = [0] * (N+1) # D[i] ; C[i] - C[i-1] (C の差分, いもす法)
for i in range(N):
if i != 0: # i=0 のときは石はもらわない
C[i] = C[i-1] + D[i]
A[i] += C[i]
stone_give = min(A[i], N-1-i) # 渡す石の総数
A[i] -= stone_give # A[i] の残りの石の数
D[i+1] += 1 # D[] を 1 増やす
D[min(N, i + stone_give + 1)] -= 1 # stone_give+1 だけ先の D[] を 1 減らす
print(*A)