K個のソート済みリストをマージ：Pythonでのヒープと分割統治ソリューション

Merge k Sorted Lists（LeetCode 23）は、k個のソート済み連結リストを1つのソート済みリストにマージする問題です。よく出題される理由は、力任せの解法が明らかで、面白い部分が「より良い2つの方法のどちらを選ぶか」にあるからです。min-heapを使う方法と分割統治法の2つがあります。どちらもO(N log k)で動作します。ここでNはノードの総数です。

問題

昇順にソートされたk個の連結リストからなる配列が与えられたとき、それらを1つのソート済み連結リストにマージして返してください。

例：[[1,4,5],[1,3,4],[2,6]] は [1,1,2,3,4,4,5,6] になります。

面接官に確認すべきこと：

• 配列が空だったり、空のリストを含んだりすることはありますか？（はい、両方に対応します。）
• 既存のノードを再利用すべきですか、新しいノードを作るべきですか？（再利用で問題なく、メモリを節約できます。）
• 重複した値はあり得ますか？（はい、なので必要ならマージを安定に保ちます。）

力任せの解法

すべての値を配列に集めてソートし、リストを再構築します。時間はO(N log N)、空間はO(N)です。動作はしますが、各リストがすでにソートされていることを無視しているので、最終的な回答として認められることはほとんどありません。

アプローチ1：min-heap

各リストに1つずつポインタを保持します。出力の次のノードは常に現在の先頭の中で最小なので、先頭をmin-heapに入れておきます。

1. 空でない各リストの先頭をヒープにプッシュします。
2. 最小のものをポップして結果に追加し、次のノードがあればプッシュします。
3. ヒープが空になるまで繰り返します。

import heapq

class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next

def merge_k_lists(lists):
    heap = []
    for i, node in enumerate(lists):
        if node:
            heapq.heappush(heap, (node.val, i, node))
    dummy = tail = ListNode()
    while heap:
        _, i, node = heapq.heappop(heap)
        tail.next = node
        tail = node
        if node.next:
            heapq.heappush(heap, (node.next.val, i, node.next))
    return dummy.next

インデックス i はタイブレーカーです。これがないと、値が等しい2つの要素でPythonが ListNode オブジェクトを比較しようとして TypeError が発生します。

計算量：N個のノードはそれぞれヒープに1回入り1回出て、各操作のコストはO(log k)なので、時間はO(N log k)です。ヒープが保持するのは最大k個のノードなので、追加の空間はO(k)です。

アプローチ2：分割統治法

ソート済みリストを2つマージするのは簡単で、O(n)時間かかります。リストをペアでマージし、その結果をまたペアでマージする、ということをトーナメントのように繰り返します。

def merge_two(a, b):
    dummy = tail = ListNode()
    while a and b:
        if a.val <= b.val:
            tail.next, a = a, a.next
        else:
            tail.next, b = b, b.next
        tail = tail.next
    tail.next = a or b
    return dummy.next

def merge_k_lists_dc(lists):
    if not lists:
        return None
    while len(lists) > 1:
        merged = []
        for i in range(0, len(lists), 2):
            other = lists[i + 1] if i + 1 < len(lists) else None
            merged.append(merge_two(lists[i], other))
        lists = merged
    return lists[0]

計算量：ラウンド数はおよそlog kで、各ラウンドですべてのN個のノードに1回ずつ触れるので、時間はO(N log k)です。ループ版では、先頭のリスト以外に必要な追加空間はO(1)だけです。

なぜ1つずつマージしないのか？

リスト1とリスト2をマージし、その結果をリスト3とマージする、という方法は自然に見えますが、コストが高くなります。最初のリストはk回のマージを通じてコピーされ続けます。すべてのリストがn個のノードを持つ場合、総作業量はおよそn × (1 + 2 + ... + k)、つまりO(n kの2乗)です。リストをペアにすることでサイズが釣り合ったままになり、各ノードに触れるのはおよそlog k回で済みます。

2つの比較

面接官が「リストはメモリに収まらないほど大きく、ストリームとして到着する」と言った場合、ヒープを使う方法は自然に拡張できます。各ストリームの現在の先頭だけが必要だからです。

テストすべきエッジケース

面接での説明の仕方

1. 力任せの解法と、それがなぜ無駄なのかを述べます。
2. 鍵となる観察を言います：出力の次のノードは、現在の先頭の中で最小のものである。
3. ヒープを選び、書いて、タイブレーカーについて言及します。
4. 計算量を述べます：時間O(N log k)、空間O(k)。
5. O(1)空間の代替案として分割統治法を提案します。
6. 空の入力と重複を含む入力でテストします。

公開前に、上記の両方のバージョンを、空のリストを含む3,000件のランダムな入力に対してソートベースのチェックと照合してテストしました。こうした簡単なランダムチェックを習慣にすることは、練習するときに良い習慣です。

関連する練習

これが簡単だったなら、まずMerge Two Sorted Listsを試し、次に同じアイデアを使うSort ListとFind Median from Data Streamを試してください。計算量の用語についてはBig Oチートシートを、動的計画法については最長共通部分文字列ガイドを、練習方法についてはコーディング面接アシスタントガイドをご覧ください。

コードを書きながら、アプローチを声に出して説明しましょう。模擬面接の記録を見れば、コードを書く前に鍵となる観察を説明できたかがわかります。ライブの評価では、AIツールに関する雇用主のルールに従ってください。