Facebookに書いた複数のキューから優先度の合計の最小を取ってくるやつ Ver.2

test_two_queues.py

from bisect import bisect_left
from heapq import heapify, heappop, heappush
from itertools import islice
from sys import maxsize

# サンプルデータ
pq0 = [(1, "a"), (2, "b"), (3, "f")]
pq1 = [(1, "c"), (2, "d"), (3, "g")]
heapify(pq0)
heapify(pq1)

pq_frontier = []  # (評価値の合計, pq0の要素, 未使用のpq1_seenのidx)
pq0_seen = []  # pq0 のうち、一度でも見たら pq0 から pq0_seen に移す
pq1_seen = []  # pq1 のうち、一度でも見たら pq1 から pq1_seen に移す


def iter_sum():
    if len(pq0) == 0 or len(pq1) == 0:
        return

    pq1_seen.append(heappop(pq1))
    while True:
        v0 = maxsize if len(pq_frontier) == 0 else pq_frontier[0][0]
        v1 = maxsize if len(pq0) == 0 else pq0[0][0] + pq1_seen[0][0]

        if v0 == maxsize and v1 == maxsize:
            return
        elif v0 <= v1:  # pq_frontier に入っている方が小さい場合
            e1 = heappop(pq_frontier)
            idx1 = e1[2]
            yield (v0, e1[1][1], pq1_seen[idx1][1])
            idx1 += 1
            if idx1 >= len(pq1_seen) and len(pq1) > 0:
                pq1_seen.append(heappop(pq1))  # pq1 → pq1_seen に移動
            if idx1 < len(pq1_seen):  # frontier の pq1_seen のインデックスを1つ進める
                heappush(pq_frontier, (e1[1][0] + pq1_seen[idx1][0], e1[1], idx1))
        else:
            yield (v1, pq0[0][1], pq1_seen[0][1])
            if len(pq1_seen) < 2 and len(pq1) > 0:
                pq1_seen.append(heappop(pq1))  # pq1 → pq1_seenに移動
            if len(pq1_seen) >= 2:  # pq1_seen の2番目の要素との組合せ
                heappush(pq_frontier, (pq0[0][0] + pq1_seen[1][0], pq0[0], 1))
            pq0_seen.append(heappop(pq0))  # pq0 → pq0_seenに移動


def add_event(add_pq, e):
    global pq_frontier
    if add_pq == 0:
        heappush(pq_frontier, (e[0] + pq1_seen[0][0], e, 0))
    else:
        if e[0] <= pq1_seen[0][0]:
            pq1_seen.append(e)
            pq1_seen.sort()  # TODO pq1_seen はスキップリストにする

            # pq_frontier を全更新
            e_idx = bisect_left(pq1_seen, e)
            pq_frontier2 = []
            while len(pq_frontier) > 0:
                e2 = heappop(pq_frontier)
                if e2[2] <= e_idx:
                    e2 = (e2[0], e2[1], e2[2] + 1)
                heappush(pq_frontier2, e2)

                # e との組合せを全部 pq_frontier に追加
                if e2[2] >= e[0]:
                    for x in pq0_seen:
                        if x[0] > e2[1][0]:
                            break
                        heappush(pq_frontier2, (e2[0] + e[0] - e2[2] + x[0] - e2[1][0], x, e[0]))
            pq_frontier = pq_frontier2
        else:
            heappush(pq1, e)


it = iter_sum()
print(list(islice(it, 4)))
# add_event(0, (1, "e"))
add_event(1, (1, "e"))
print(list(islice(it, 20)))