ABC141-D問題のTLE

AtCoderのABCに定期的に（と言ってもまだ５，６回）参加している。

リアルタイムの参加はできなかったが，ABC141のD問題に挑戦したので記事にしようと思う。

問題は公式サイトを参考にしてください。

D - Powerful Discount Tickets

一番値段の高いものにディスカウントチケットを使っていけばいいなというのは（直感で）わかる。また，一枚ずつ何回も使うのと，複数枚を同時に使うことの差もないことも少し実験すれば確認できる（本当は証明が必要。公式editorialには証明も載っています）。

とりあえず書いたコードがこれ。

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

maxi = 0
idx = 0
for i in range(M):
   maxi = max(A)
   idx = A.index(maxi)
   A[idx] = maxi//2
print(sum(A))

なんて素直な実装なのでしょうか。

listから一番大きい要素を取り出して2で割って（ディスカウントして），その要素の値を更新するということをチケット枚数分（M回）繰り返すだけ。これで，19テストケースのうち12ケースがTLEでした。他はＡＣ。なので考え方自体は間違っていないと感じたので，高速化を検討します。正直言って知らないと無理だと思うのですが，heapqueue（ヒープキュー）というのを使えば解決します。

heapq --- ヒープキューアルゴリズム — Python 3.7.5rc1 ドキュメント

これを使うとlist内の一番小さい要素を高速に取り出すことができますよ，というやつです。基本情報技術者試験とかでスタックとかキューとかあるあれですね。それの一種。細かい内部のことはよくわかっておりませんがとりあえずこれ使えば速いらしいので使ってみました。

import heapq

N, M = map(int,input().split())
A = list(map(int,input().split()))
#大きいものから取り出したいので-1倍してqueueに入れる

Q = []
for a in A:
   heapq.heappush(Q,-a)

for i in range(M):
   first = -heapq.heappop(Q)
   first = first//2
   heapq.heappush(Q,-first)

print(-sum(Q))

ここで注意ですが，pythonのheapqはpopしたときに一番小さい値をとりだします。今は一番大きな値を取り出したいので，キューに格納するときにすべての値をあらかじめ-1倍してから格納しておきます。そうすることで大小関係が反転して，popしたときに取り出した要素が一番大きな要素です（持ちろん取り出すときにも-1倍して元の値に戻します）。あとはそれを指定回数繰り返せば問題なくACでした。