あらこま日記

よろしく

ABC141D - Powerful Discount Tickets

AtCoder Python priority queue heapq

問題

【概要】
N個の品物があり、それぞれ $A_{1}...A_{N}$ 円である。
M枚の割引券があり、一回使うと任意の品物の値段を1/2にできる(小数以下切捨、同じ品物にも重ねて使える)。
全ての品物を買えるような最小の金額を求める。
【制約】
$1\leq N,M \leq 10^{5}$

考える

一番値段の高い品物に割引券使う。これをM回やればよさそう。
割引するごとにソートして一番高い値段を得る。⇒TLE...
時間切れまでいろいろやるも終了。

解説見る

(pythonの)max()は $O(N)$ なので、愚直に繰り返すと $O(NM)$ でTLEする。
優先度付きキュー（priority queue, heap queue）を使うと、配列について以下の操作を比較的高速に行える。
- 値の追加 $O(\log_{}{N})$
- 値の削除 $O(\log_{}{N})$
- 最小値(最大値)の取得 $O(1)$
最小値(最大値)のみを得たい、のみに処理したい、を何回も繰り返す場合は優先度付きキューを使うと速い！！！
pythonだとimportして使える。

import heapq

heapq.heapify(L) #list(L)をheapに変換
heapq.heappop(L) #最小値を取り出して返す
heapq.heappush(L, a) #heapに値(a)を追加

注意！pythonのheapqは最小値が根にくるので、最大値が欲しいなら全部の要素に-1をかけてからheapに変換する。

これらを使うと、「一番高い品物を選んで半額にする」という操作を $O(\log N)$ でできる。よって、全体で $O(M \log N)$ で間に合う。

ACコード(python)

import heapq

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

#全要素に-1かける
for i in range(len(A)):
    A[i] *= -1

heapq.heapify(A)
for _ in range(M):
    x = heapq.heappop(A)
    heapq.heappush(A, (-x//2)*(-1)) #切り下げのためマイナスを2回かける

print(-sum(A))#全要素を負にしたので戻すためにマイナス付ける

おわりに

くぇうぇ笑

参考(ありがとうございます…！)