ABC141 D問題 Powerful Discount Tickets を優先度付きキューで解く

ABC141D問題を解説します。
優先度付きキュー(heapq)の練習で取り組みました。

問題の概要

$N$ 個の商品と $M$ 枚の割引券があります。
1つの商品に複数回割引券を使うことができます。
全部の商品を買うために必要な最小の金額を求める問題です。

難しいポイント

方針はすぐに思いつくかもしれませんが、普通の配列を使って最大値を取り出したり、要素を追加したりすると、 $O(N)$ の計算量がかかってしまいます。
$N$ 、 $M$ は最大で $10^5$ なので、タイムオーバーします。

解き方の方針

一番高い値段に割引券を使うのが一番お得になります。
1枚ずつ1番高い値段の商品に割引券を使っていきます。
最大値を取り出す、要素の削除、要素の追加を高速に処理するために優先度付きキュー(heapq)を使います。
heapqのheappopで $O(1)$ で最小値を取り出せますし、 $\log{N}$ で要素を追加できます。

実装

heapq.heappopは最小値を取り出すので、データを読み込むときに-1をかけて、大小を逆転させておきます。

切り捨てのときにマイナスになっていると、値が大きくなってしまいます。
例えば、-13を2で割った切り捨ては-7になってしまいます。
なので、2で割る前に-1をかけて、割ってから-1をかけてあげます。

from heapq import heapify, heappop, heappush
N, M = map(int, input().split())
# 全ての値に-1をかけて大小を逆転させる
a_arr = list(map(lambda x: int(x) * -1, input().split()))
# 配列を優先度付きキューに変換する
heapify(a_arr)

for i in range(M):
    most_expensive = heappop(a_arr)
    # マイナスのまま切り捨ての割り算をすると大きい値になるので、
    # 一度-1をかけてから2で割って、最後に-1をかける
    most_expensive = ((-1 * most_expensive) // 2) * -1
    heappush(a_arr, most_expensive)

print(-1 * sum(a_arr))