PythonでK個のソート済みリストを1つにマージする方法（ヒープ活用の実装例）

はじめに

複数のソート済みリストが与えられたとき、それらを1つのソート済みリストにまとめる問題は、アルゴリズム学習における定番の課題です。この問題を効率的に解くには、ヒープ（優先度付きキュー）というデータ構造を利用するのが有効です。

たとえば、[1, 4, 5]、[1, 3, 4]、[2, 6] という3つのソート済みリストがある場合、マージ後の最終的なリストは [1, 1, 2, 3, 4, 4, 5, 6] になります。

アルゴリズムの手順

この問題は、次の手順で解くことができます。

1. リスト群 lists のサイズを n とします。
2. 空のヒープ（heap）を用意します。
3. 各リスト lists[i] を走査し、リストが空でなければその先頭要素を (値, 行番号, 列番号0) のタプルとしてヒープに挿入します。
4. 結果を格納するための空リスト res を用意します。
5. ヒープが空になるまで、以下の処理を繰り返します。
   • ヒープから最小の要素 (num, row, col) を取り出し、num を res の末尾に追加します。
   • col が lists[row] のサイズ − 1 より小さい場合、同じ行の次の要素 lists[row][col + 1] を (値, 行番号, 列番号+1) としてヒープに挿入します。
6. すべての処理が終わったら res を返します。

ポイントは、常に「各リストのまだ消費していない先頭要素」だけをヒープに保持することです。これにより、全体を無理なく昇順に取り出せる仕組みになっています。

Pythonでの実装例

理解を深めるために、実際のコードを見てみましょう。標準ライブラリの heapq モジュールを使用します。

import heapq

class Solution:
    def solve(self, lists):
        heap = []
        # 各リストの先頭要素をヒープに登録
        for i, row in enumerate(lists):
            if row:
                heapq.heappush(heap, (row[0], i, 0))

        res = []
        while heap:
            num, row, col = heapq.heappop(heap)
            res.append(num)
            # 同じ行に次の要素が残っていればヒープへ追加
            if col < len(lists[row]) - 1:
                heapq.heappush(heap, (lists[row][col + 1], row, col + 1))
        return res

ob = Solution()
lists = [[], [], [11, 13], [], [4, 4, 14], [4], [11], [1, 8]]
print(ob.solve(lists))

なお、このコードでは空のリストも混在した入力を正しく扱える点が特徴です。空リストは「if row」の判定によって自動的にスキップされます。

入力

[[], [], [11, 13], [], [4, 4, 14], [4], [11], [1, 8]]

出力

[1, 4, 4, 4, 8, 11, 11, 13, 14]

計算量について

全要素数を N、リストの個数を K とすると、このアルゴリズムの時間計算量は O(N log K) です。各要素は最大でも1回ずつヒープに挿入・削除され、ヒープには常に高々 K 個の要素しか存在しないためです。単純にすべての要素を連結してソートする方法（O(N log N)）よりも、K が小さい場合に有利になります。また、空間計算量はヒープに保持する要素数分の O(K) です。

まとめ

K個のソート済みリストのマージは、ヒープを使うことで効率的かつシンプルに実装できます。heapq の heappush / heappop を組み合わせるだけで、追加のライブラリ不要で動作するのも魅力です。リンクトリスト版への応用や、分割統治法によるマージとの比較などにも挑戦してみてください。