【Python】ヒープを使ってk個のソート済みリストを効率的にマージする方法

複数のソート済みリストが与えられたとき、それらをすべてマージして1つのソート済みリストを作ることを考えます。この問題は、ヒープ（優先度付きキュー）というデータ構造を使うことで効率的に解くことができます。

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

アルゴリズムの手順

以下の手順に従って問題を解きます。

• まず、空の最小ヒープを1つ作成します。
• lists 内の各連結リスト l に対して以下を繰り返します。
  • l が空（None）でない場合は、その先頭ノードをヒープに挿入します。
• 結果リストの先頭ポインタ res を null に、末尾を追跡するためのポインタ res_next も null に初期化します。
• 無限ループを実行します。
  • temp := ヒープから最小値を持つノードを取り出します。
  • ヒープが空になった場合は、res を返して終了します。
  • res がまだ null の場合（最初のノードを接続するとき）：
    • res := temp、res_next := temp とします。
    • temp を次のノードへ進めます。
    • temp が null でなければ、temp をヒープに挿入します。
    • res.next を null に設定します。
  • それ以外の場合：
    • res_next.next := temp とし、temp を次のノードへ進め、res_next も次へ進めます。
    • temp が null でなければ、temp をヒープに挿入します。
    • res_next.next を null に設定します。

計算量の目安

全ノード数を N、リストの本数を k とすると、各ノードは高々1回ずつヒープへの挿入と取り出しを行うため、時間計算量は O(N log k) になります。また、ヒープには常に最大 k 個のノードしか保持しないため、空間計算量は O(k) です。すべての要素を連結してからソートする方法（O(N log N)）と比べ、リストの本数 k が小さい場合に有利なアプローチです。

実装例

以下のPython実装を見ると、より理解が深まるでしょう。

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

def make_list(elements):
   head = ListNode(elements[0])
   for element in elements[1:]:
      ptr = head
      while ptr.next:
         ptr = ptr.next
      ptr.next = ListNode(element)
   return head

def print_list(head):
   ptr = head
   print('[', end = "")
   while ptr:
      print(ptr.val, end = ", ")
      ptr = ptr.next
   print(']')

class Heap:
   def __init__(self):
      self.arr = []

   def print_heap(self):
      res = " "
      for i in self.arr:
         res += str(i.val) + " "
      print(res)

   def getVal(self,i):
      return self.arr[i].val

   def parent(self,i):
      return (i-1)//2

   def left(self,i):
      return 2*i + 1

   def right(self,i):
      return 2*i + 2

   def insert(self,value):
      self.arr.append(value)
      n = len(self.arr)-1
      i = n
      while i != 0 and self.arr[i].val<self.arr[self.parent(i)].val:
         self.arr[i],self.arr[self.parent(i)] = self.arr[self.parent(i)],self.arr[i]
         i = self.parent(i)

   def heapify(self,i):
      left = self.left(i)
      right = self.right(i)
      smallest = i
      n= len(self.arr)
      if left<n and self.getVal(left)<self.getVal(smallest): smallest = left
      if right <n and self.getVal(right)<self.getVal(smallest): smallest = right
      if smallest!=i:
         self.arr[i],self.arr[smallest] = self.arr[smallest],self.arr[i]
         self.heapify(smallest)

   def extractMin(self):
      n = len(self.arr)
      if n==0:
         return '#'
      if n== 1:
         temp =self.arr[0]
         self.arr.pop()
         return temp
      root = self.arr[0]
      self.arr[0] = self.arr[-1]
      self.arr.pop()
      self.heapify(0)
      return root

class Solution(object):
   def mergeKLists(self, lists):
      heap = Heap()
      for i in lists:
         if i:
            heap.insert(i)
      res = None
      res_next = None
      while True:
         temp = heap.extractMin()
         if temp == "#":
            return res
         if not res:
            res = temp
            res_next = temp
            temp = temp.next
            if temp:
               heap.insert(temp)
            res.next = None
         else:
            res_next.next = temp
            temp = temp.next
            res_next=res_next.next
            if temp:
               heap.insert(temp)
            res_next.next = None

ob = Solution()
lists = [[1,4,5],[1,3,4],[2,6]]
lls = []
for ll in lists:
   l = make_list(ll)
   lls.append(l)
print_list(ob.mergeKLists(lls))

入力

[[1,4,5],[1,3,4],[2,6]]

出力

[1, 1, 2, 3, 4, 4, 5, 6]