はじめに
本記事はAtCoder Beginner Contest 141 D問題の解法を含みます．
解説を見ずに自力で解くつもりの方はご注意ください．
前書き
ABC141のD問題「Powerful Discount Tickets」を解きました．
問題を見ると，「値段Aiのリストの中から最大値を取得してその品物に割引券を1枚適用する(半額にする)」という処理をM回繰り返せば最小の購入金額が求まりそうですが，そのまま愚直に実装したらTLEになってしまいました．
N個の品物の中から最大値を取得するにはO(N)の計算量がかかるため，これをM回繰り返すと時間計算量はO(NM)になります．この問題ではN,M<=10**5という制約があるため，MとNの値によっては実行時間を超過してしまいます．
このような問題では優先度付きキューを使用すると良いです．
優先度付きキュー
優先度付きキューとは，キューに対して要素を優先度付きで追加するデータ型で，ヒープ構造を持っています．
一般的な配列と比較して要素の参照が高速であることが特徴で，リスト内のデータの操作と最大値や最小値の取得を繰り返し行いたい場合に便利です．
優先度付きキューにおけるそれぞれの処理の計算量は以下の通りです．
| 処理 | 計算量 |
|---|---|
| リストのヒープ化 | O(N) |
| 要素の追加 | O(logN) |
| 要素の削除 | O(logN) |
| 最小値の取得 | O(1) |
heapq
Pythonでは標準ライブラリであるheapqモジュールが優先度付きキューのアルゴリズムを提供しているので，
import heapq
で使用できます．
heapqのよく使うメソッドは以下の3つです．
heapq.heapify(a) #リストaを優先度付きキューに変換(ヒープ化)する
heapq.heappush(a,x) #ヒープ化されたリストaに要素xをプッシュする
heapq.heappop(a) #ヒープ化されたリストaから最小の要素をポップする
heapqの使用例を以下に示します．
import heapq
a = [2, 13, 8]
heapq.heapify(a) #リストのヒープ化
print(a)
heapq.heappush(a,5) #プッシュ
print(a)
print(heapq.heappop(a)) #最小値の取得と表示
print(a)
[2, 13, 8]
[2, 5, 8, 13]
2
[5, 13, 8]
heapqで最大値を取得する
heapqでは最小値の取得しかできません．
heapqで最大値を取得したい場合は，全要素を-1倍してから最小値を取得し，再度-1倍する必要があります．
import heapq
a = [2, 13, 8]
print(a)
a = list(map(lambda x:x*(-1), a)) #各要素を-1倍
heapq.heapify(a) #リストのヒープ化
print(a)
print(heapq.heappop(a)*(-1)) #最大値の取得と表示
print(a)
[2, 13, 8]
[-13, -2, -8]
13
[-8, -2]
ABC141のD問題の解答例
冒頭で述べたABC141のD問題はheapqを使用することでTLEすることなく解けます．
import heapq
N,M = map(int, input().split())
A = list(map(lambda x:int(x)*(-1), input().split())) #-1倍してからリストに格納
heapq.heapify(A) #優先度付きキューに変換
for _ in range(M):
max_value = heapq.heappop(A)*(-1) #最大値の取得
heapq.heappush(A, (max_value//2)*(-1)) #半額にして-1倍してからキューに戻す
print((-1)*sum(A)) #最後に-1倍を忘れずに
おわりに
閲覧ありがとうございました．
AtCoderではしばしば複数回のソートや要素の操作と最大値or最小値の取得が要求される問題があり，リストを使った実装ではTLEしてしまう問題でも優先度付きキューを使えばACできる場合があります．