Help us understand the problem. What is going on with this article?

優先度付きキュー(ヒープ)の使い方

概要

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

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

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

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

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

毎回ソートする場合(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))

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

成功例
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】切り捨て除算演算子を使った切り上げ【算数】

Why do not you register as a user and use Qiita more conveniently?
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away
Comments
Sign up for free and join this conversation.
If you already have a Qiita account
Why do not you register as a user and use Qiita more conveniently?
You need to log in to use this function. Qiita can be used more conveniently after logging in.
You seem to be reading articles frequently this month. Qiita can be used more conveniently after logging in.
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away