競プロ(Python)で優先度キューはheapq？queue.PriorityQueue？どちらを使うべき？

Pythonの優先度付きQueueは以下の２種類があります。速度を比べてみました。

• heapq
• queue.PriorityQueue

はじめに

一見、

class queue.PriorityQueue(maxsize=0)
優先順位付きキューのコンストラクタです。maxsize はキューに置くことのできる要素数の上限を設定する整数です。いったんこの大きさに達したら、挿入はキューの要素が消費されるまでブロックされます。もし maxsize が0以下であるならば、キューの大きさは無限です。

ということでqueue.PriorityQueueは名前からしても良さそうです。しかし、queue自身は、

queue モジュールは、複数プロデューサ-複数コンシューマ(multi-producer, multi-consumer)キューを実装します。これは、複数のスレッドの間で情報を安全に交換しなければならないときのマルチスレッドプログラミングで特に有益です。

ということで、明らかに重い(同期のための動作が入っていそう)です。

やったこと

• $ 2 \cdot 10^{5} $ 個のランダムな正の整数($\leq 10^{9}$) からなる配列aを用意します
• 別の$ 2 \cdot 10^{5} $ 個のランダムな正の整数($\leq 10^{9}$) からなる配列bを用意します

次に、以下の動作を$ 2 \cdot 10^{5} $回行います。(i回目の操作とします)

• aから最も小さい数字をとる。次にbのi個目の要素をaに追加する

これを、heapqとqueueで比較します。なお、上記の作業を行うとき、heapqはpopしてpushするより、pushpopを使った方が高速なのでその速度も比較します。

結果

--- Heapq, poppush
elapsed_time:0.15400314331054688[sec]
--- Heapq, pop then push
elapsed_time:0.20299005508422852[sec]
--- Queue
    Step1: create init queue
elapsed_time:0.37195301055908203[sec]
    Step2: pop then push
elapsed_time:0.873100996017456[sec]

ということでPythonで(競技プログラミングにおいて)優先度付きキューを使うときはheapqを用いましょう。尚、queue.PriorityQueueはinitがないので最初に配列aを１つずつ作ったqにpushしています。

コード

import queue
import heapq

import time
import random
from copy import deepcopy
n = 2 * 10**5

inqueue = [random.randint(0, 10**9) for _ in range(n)]
inheapq = deepcopy(inqueue)

data = [random.randint(0, 10**9) for _ in range(n)]

resqueue = []
resheapq = []

print("--- Heapq, poppush")
start = time.time()
heapq.heapify(inheapq)
for i in range(n):
    x = heapq.heappushpop(inheapq, data[i])
    resheapq.append(x)
elapsed_time = time.time() - start
print("elapsed_time:{0}".format(elapsed_time) + "[sec]")

print("--- Heapq, pop then push")
start = time.time()
heapq.heapify(inheapq)
for i in range(n):
    x = heapq.heappop(inheapq)
    resheapq.append(x)
    heapq.heappush(inheapq, data[i])
elapsed_time = time.time() - start
print("elapsed_time:{0}".format(elapsed_time) + "[sec]")

print("--- Queue")
print("    Step1: create init queue")
start = time.time()
q = queue.PriorityQueue(n)
start = time.time()
for i in range(n):
    q.put(inqueue[i])
elapsed_time = time.time() - start
print("elapsed_time:{0}".format(elapsed_time) + "[sec]")
print("    Step2: pop then push")
start2 = time.time()
for i in range(n):
    x = q.get()
    resqueue.append(x)
    q.put(data[i])
elapsed_time = time.time() - start2
print("elapsed_time:{0}".format(elapsed_time) + "[sec]")