ABC320 E - Somen Nagashi (Diff: 1096 緑) の解説 [AtCoder] [Python]

問題へのリンク

https://atcoder.jp/contests/abc320/tasks/abc320_e

問題概要

先頭から順に1~Nの番号がついたN人が一列に並んでいる。これからM回そうめんが流される。

そうめん i は時刻 T_i に流され、量は W_iである。そうめん i が流されたとき、列の先頭にいる人が全てを得る。その人は列から外れ、時刻T_i + S_iに元いた場所に戻ってくる。

それぞれの人がどれだけそうめんを得られたか求めよ。

制約

  • 1 ≦ N, M ≦ 105
  • 0 < T_1 < T_2 … <T_M ≦ 109
  • 1 ≦ S_i, W_i ≦ 109

解説

問題文に出てくる「列」を「そうめん列」と呼ぶことにする。

i番目のそうめんを食べることができるのは、そうめん列の中で番号が最小の人である。この部分は「優先度付きキュー」(heapq)で高速に実装できそうだ。

また、ある人がそうめん i を食べた時、戻ってくる時間はT_i + S_i だと分かっている。そのため、そうめん列から外れた人を管理する「待機列」を作り、そうめん列に戻る時間が小さい順に並べる。そして、そうめん i が配られる前の各タイミングで、「 最も順番が早い人が戻る時間が T_i 以下の時、そうめん列に戻す」操作を繰り返す。

この「待機列」も、(戻る時間, 人の番号)のタプルを要素にすることで、優先度付きキューで実装できそうだ。

実装(Python)

from heapq import *

N, M = map(int, input().split())

# そうめん列, heapq化
somen_que = list(range(N))
heapify(somen_que)

# そうめんを得た量を管理するリスト
somen_counter = [0] * N

# 待機列
wait_que = []

for _ in range(M):
    T, W, S = map(int, input().split())

    # もし待機列に人がいて,
    # 先頭の戻る時間がT以下なら戻す
    while wait_que and wait_que[0][0] <= T:
        t, p = heappop(wait_que)
        heappush(somen_que, p)

    # そうめん列に誰もいないなら何もしない
    if not somen_que:
        continue

    # p: そうめん列の先頭の人
    p = heappop(somen_que)
    somen_counter[p] += W

    # (pが戻る時間, 番号)のタプルを追加
    heappush(wait_que, (T + S, p))

# 答えを出力
for i in somen_counter:
    print(i)