Pythonで優先度付きキューを実装してみる-ABC141

今回は優先度付きキューをPythonのheapqで実装していきます。

優先度付きキューとは

優先度付きキューとは一言で説明すると最大値、最小値を普通のリストよりも早く見つけることができるアルゴリズムです。

通常のリストでは最大値、最小値を見つける計算量が(O(N))であるのにたいして、優先度付きキューでは(O(logN)))で処理することが可能です。

そしてPythonでは優先度付きキューを実装するためにheapqモジュールが用意されています。

実際に問題を解く前に基本的なheapqモジュールの使い方を解説していきます。

heapqのメソッド

以下の4つのheapqのメソッドを覚えておきましょう。

• heapq.heappush(heap,item) #要素の挿入
• heapq.heapppop(heap) #最小値の取り出し
• heapq.heappushpop(heap,item) #最小値を取り出したあとに要素を挿入(heappushとheappopを別々に実行するよりも効率的に実行できる)
• heapq.heapify(x) #リストを優先度付きキューに変換

import heapq

ls = [9,2,1,4,0,-3]
heapq.heapify(ls)  
print(ls)

print(heapq.heappop(ls))  
print(ls)

heapq.heappush(ls, 5)  
print(ls)

print(heapq.heappushpop(ls, 7)) 
print(ls)

出力

[-3, 0, 1, 4, 2, 9]
-3
[0, 2, 1, 4, 9]
[0, 2, 1, 4, 9, 5]
0
[1, 2, 5, 4, 9, 7]

最大値を取り出す

heapqはそのままでは最小値しか取り出すことができません。しかし、その性質を利用して最大値を取り出すことが出来ます。

最大値を取り出すためには全ての要素に-1をかけることで取り出すことが可能です。(-1をかけることによって最大値が最小値となるため)

import heapq

ls = [9,2,1,4,0,-3]
ls =[-1*i for i in ls]

heapq.heapify(ls) 

print(heapq.heappop(ls)*(-1)) #最大値の取り出し
print(ls)

print(heapq.heappushpop(ls, 7)*(-1)) # 最大値を取り出したあとに要素を挿入
print(ls)

出力

9
[-4, -2, -1, 3, 0]
4
[-2, 0, -1, 3, 7]

例題(ABC141-D)

それでは次は実際に問題を解いていきたいと思います。

今回は以下の問題を扱います。
atcoder.jp

問題文

高橋くんは N 個の品物を 1個ずつ順番に買う予定です。
i番目に買う品物の値段は Ai円です。
高橋くんは M枚の割引券を持っています。
品物を買う際に割引券を好きな枚数使うことができます。
X円の品物を買う際に Y 枚の割引券を使った場合、その品物を X/2^Y円(小数点以下切り捨て)で買うことができます。
最小で何円あれば全ての品物を買うことができるでしょうか。

制約

入力は全て整数
1≤N,M≤105
1≤Ai≤109

実装

この問題は優先度付きキューさえ知っていれば非常に簡単に解くことができます。

考え方としては、最も値段の高いものから順に割引券を使用すればいいだけです。

ただし、単純にmaxを使ったり、sortを使ってしまっては実行が間に合いません。

そこで優先度付きキューの出番です。今回は以下のように実装してみました。

import heapq
n,m = map(int,input().split())
a = list(map(lambda x: -int(x),input().split()))

heapq.heapify(a)

for i in range(m):
    a_max = heapq.heappop(a)
    heapq.heappush(a,-(-a_max//2))
print(-sum(a))

このコードでは入力を受け取る時に値がマイナスになるようにしています。そうすることで受け取ってからマイナスにするよりも高速に処理できます。

そしてあとは先ほどやったことと同じです。
リストを優先度付きキューに変換し、割引券の数だけ最大値を取り出したらそれを半額(a_max//2だと計算誤差がでてしまうので-(-a_max//2)という書き方にしました)にしてpushするという操作を繰り返します。

ヒープソートの実装

おまけですが、heapqを利用することで以下のようにヒープソートも実装することができます。

def heapsort(ls):
    h = []
    for value in ls:
        heapq.heappush(h, value)
    return [heapq.heappop(h) for i in range(len(h))]

heapsort([1,4,0,3,-1,-3])

出力

[-3, -1, 0, 1, 3, 4]

さいごに

優先度付きキューをPythonで実装するには基本的なheapqの使い方さえ覚えれば簡単に実装することができます。

最大値、もしくは最小値を取り出す問題で処理速度をあげたい時は優先度付きキューを使ってみましょう。

おわり