優先度付きキュー（ヒープ）の使い方 #Python

概要

AtCoderのABC141で優先度付きキュー（ヒープ）を用いる問題に出くわしたので、その備忘録です。公式の解説があるため詳しい説明は省きますが、この問題を解くに当たって下記の処理を行う必要が出てきます。

最も高い値段を持つ品物に対して割引（50%オフ）を適用するという処理を、割引券がなくなるまで繰り返す

単純に考えると、これは下記の処理を繰り返せば実現できそうです。

値段の配列をソートする
最も値段の高い要素を取得し、半額にした上で配列に戻す

ただしこの方法の場合、毎回ソートしているため計算量が増大し、与えられた制限時間以内に処理が終わりません。そこで最大値の取得を効率的に繰り返すため、優先度付きキュー（ヒープ）を用いる必要があります。下記にソートを用いた失敗例も含めて記載します。

毎回ソートする場合（TLE）

失敗例

n, m = map(int, input().split())
a = []

for i in input().split():
    a.append(i)

# 割引券がなくなるまで繰り返す
for j in range(m):
    # 最大値を取得するため、毎回ソートを行う
    a.sort()
    # 最も高い値段を取り出す
    highest = a[-1]
    # 半額にしてリストに戻す
    a[-1] = highest // 2

print(sum(a))

優先度付きキューを用いる場合（AC）

成功例


import heapq

n, m = map(int, input().split())
a = []

for i in input().split():
    # Pythonのheapは最小値がpopされるため、マイナスを付与した上でheapに入れる
    heapq.heappush(a, -int(i))

# 割引券がなくなるまで繰り返す
for j in range(m):
    # 最も高い値段を取り出す
    highest = heapq.heappop(a)
    # 半額にしてheapに戻す
    heapq.heappush(a, -(-highest // 2))

print(-sum(a))

注意点としては下記になります。

Pythonの優先度付きキューはpopの際に最小値を返すため、あらかじめ負の値にした上でpushする（最後にsumを出力する際に再度マイナスを付けて正数に戻す）
負の数に対して//を用いると、負の方向に切り上がる（絶対値が増加する）ため、一旦正数に戻してから切り下げる

そもそも何故Heapを用いると処理が速くなるのかに関しては、また別の記事で紹介したいと思います。

参考

heapq ヒープキューアルゴリズム
 ヒープをわかりやすく解説してみた
 【Python】切り捨て除算演算子を使った切り上げ【算数】