Pythonで2つのソート済みリストをマージして1つのソート済みリストを作成する方法

2つのソート済みリスト A と B が与えられたとき、それらをマージして、1つのソート済みリスト C を作成することを考えます。なお、2つのリストのサイズは異なっていても構いません。

例えば、A = [1, 2, 4, 7]、B = [1, 3, 4, 5, 6, 8] の場合、マージ後のリスト C は [1, 1, 2, 3, 4, 4, 5, 6, 7, 8] となります。

アルゴリズムの考え方

この問題は、2つのリストを先頭から順に比較しながら要素を取り出していく「マージ処理」で解くことができます。各リストに対してポインタ（インデックス）を1つずつ用意し、小さい方の要素を結果のリストに追加していきます。手順は以下の通りです。

• 結果を格納するための新しいリスト x を用意する
• インデックス i = 0、j = 0 と初期化する
• i が lst0 のサイズ未満 かつ j が lst1 のサイズ未満である間、以下を繰り返す
  • lst0[i] > lst1[j] の場合：x の末尾に lst1[j] を追加し、j を +1 する
  • lst0[i] < lst1[j] の場合：x の末尾に lst0[i] を追加し、i を +1 する
  • それ以外（等しい場合）：x の末尾に lst0[i] と lst1[j] の両方を追加し、i と j をそれぞれ +1 する
• ループ終了後、lst0 に残った要素があればすべて x に追加する
• 同様に、lst1 に残った要素があればすべて x に追加する
• x を返す

どちらか一方のリストの要素を使い切った後は、もう片方のリストに残っている要素はすでにソートされているため、そのまま順番に追加するだけで正しい結果が得られます。

実装例

それでは、実際のコードを見て理解を深めましょう。

class Solution:
    def solve(self, lst0, lst1):
        x = []
        i = 0
        j = 0
        while(i < len(lst0) and j < len(lst1)):
            if(lst0[i] > lst1[j]):
                x.append(lst1[j])
                j = j + 1
            elif(lst0[i] < lst1[j]):
                x.append(lst0[i])
                i = i + 1
            else:
                x.append(lst0[i])
                x.append(lst1[j])
                i = i + 1
                j = j + 1
        while(i < len(lst0)):
            x.append(lst0[i])
            i = i + 1
        while(j < len(lst1)):
            x.append(lst1[j])
            j = j + 1
        return x

ob = Solution()
print(ob.solve([1, 2, 4, 7], [1, 3, 4, 5, 6, 8]))

入力

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

出力

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

計算量について

このアルゴリズムの時間計算量は O(m + n) です。ここで m と n はそれぞれ lst0 と lst1 の要素数です。各要素は必ず一度だけ比較・追加されるため、非常に効率的です。また、空間計算量も結果リストの分だけの O(m + n) となります。

なお、Pythonでは標準ライブラリの heapq.merge() を使うことでも、ソート済みイテラブルを効率的にマージできます。ただし、マージソートの仕組みを理解する上で、今回のような手動での実装は非常に良い学習材料になります。