atcoder ABC320 E - Somen Nagashi のSortedListを使った実装

atcoder.jp

・参考　
トヨタ自動車プログラミングコンテスト2023#5(AtCoder Beginner Contest 320) - YouTube
Sorted List — Sorted Containers 2.4.0 documentation

・説明
　SortedListを使うことで先頭の要素(昇順になっているので最小値)をpop(0)を使い、O(log (N))で取得できる。普通のリストではpopに引数を付けるとO(N)になり、TLEになる。他にもheapqやdequeを使うことで解くことができるがSortedListが楽だと思う。
　プログラムの説明は今、列にいる人をorder、列から外れた人をt_whoで管理して、T以下の時間のものをorderに戻す。詳しい説明はコメントアウトしている。

from sortedcontainers import SortedList

n,m=map(int,input().split())

ans=[0 for _ in range(n)]
order=SortedList([i for i in range(n)])  #列にいる人の順番.
t_who=SortedList()   #(時間,何番目の人)を追加していく、何番目の人がどの時間に列に戻るか表す。

for _ in range(m):
    t,w,s=map(int,input().split())
    
    while len(t_who)!=0:
        tt,who=t_who.pop(0)
        #t_whoからt以下の時間の人を列orderに戻す.
        if tt<=t:
            order.add(who)
        #そうでなかったら、popしたものをもとに戻してbreak.
        else:
            t_who.add((tt,who))
            break
    
    if len(order)==0:
        continue
    p=order.pop(0)   #そうめんを食べる人.
    ans[p]+=w
    t_who.add((t+s,p))

for i in range(n):
    print(ans[i])