Merge k Sorted Lists
listsの各行として、k個の整数リストが与えられます。各行は非減少順にソートされており、行ごとに長さが異なる場合があり、空の行はありません。
すべての行のすべての値を含む1つのリストにマージし、非減少順にソートして返してください。1つの行または複数の行に同じ値が複数回現れる場合、結果にもその回数分だけ現れます。
関数
- listsinteger-2d-array
- 行ごとに並べ替えられたリスト。長さは異なる場合があります
- 戻り値integer-array
- すべての行のすべての値を、1つのソート済みリストに
制約
1 ≤ lists.length ≤ 1041 ≤ lists[i].lengthであり、すべての行を合わせても最大で104個の値を含みます-104 ≤ lists[i][j] ≤ 104- 各行は非減少順に並べ替えられています。
例
- 入力
- lists = [[2, 6, 9], [1, 4, 10], [3, 5]]
- 出力
- [1, 2, 3, 4, 5, 6, 9, 10]
- 説明
- 行は2、1、3で始まるので、最初は1です。次に先頭の数字は2、4、3なので、続いて2となり、以降も同様です。3行目は5の後で終わるため、6、9、10が残ります。
- 入力
- lists = [[5], [-2, 5, 7], [0, 5]]
- 出力
- [-2, 0, 5, 5, 5, 7]
- 説明
- 3つの5はそれぞれ異なる行から来ており、3つとも残ります。負の値である
-2は0より前に並びます。
- 入力
- lists = [[4, 8]]
- 出力
- [4, 8]
- 説明
- 1 行だけの場合、結合するものはありません。行はすでにソートされているため、それが答えです。
発展問題
各行の値を1つずつ含む最小の範囲[a, b]を見つけてください。各行の先頭要素を格納したヒープと最大の先頭要素を使えば、O(N log k)で実現できますか？
ヒント
1つずつ開いてください。開くたびに少しずつ答えに近づきます。
各行はソートされています。すべての値の中で最も小さい可能性があるのはどの値でしょうか？
答えの次の値は、常に各行の未使用の値のうち最小のものです。それを取り出すと、変化する値はそのうちの1つだけです。
各行の最初の未使用値を、行番号を付けて min-heap に格納します。最小値を取り出して追加し、次の値が同じ行にあれば、それを追加で格納します。
解説
各行はソートされているため、未使用の値のうち最小のものは、常にいずれかの行の未使用の値の先頭、つまりその行の先頭要素です。すべてのk個の先頭要素を走査すると、N個の値それぞれについてkステップかかります。先頭要素を格納した最小ヒープを使えば、最小の値をO(log k)で取り出せるため、全体の計算量をO(N·k)からO(N log k)に減らせます。
すべての値について、k 個のヘッドを比較する
正しいが、最大のテストでは終わらない
考え方
各行の先頭にインデックス pos[r] を保持します。各ラウンドではすべての先頭要素を読み取り、最小値を追加して、その行のインデックスを進めます。これはマージソートのマージ処理を、2つのリストから k 個のリストに拡張したものです。最初の例では、先頭要素は2、1、3なので、1が取り出され、2行目の先頭要素は4になります。
N 個の値それぞれについて、k 行すべてを調べる必要があります。各行に値が1つずつある 10^4 行の場合、行のチェック回数は 10^8 回になります。
アルゴリズム
- すべての行について
pos[r] = 0に設定し、値の数をNとして数えます。 N回繰り返します。pos[r]がまだ行の範囲内にあるすべての行を確認し、先頭の値が最小の行を覚えておきます。- その先頭の値を結果に追加し、その行の
posに 1 を加えます。 - 結果を返します。
def mergeKLists(lists):
pos = [0] * len(lists) # pos[r]: index of the first unused value in row r
total = sum(len(row) for row in lists)
merged = []
for _ in range(total):
best = -1 # the row whose head is the smallest so far
for r in range(len(lists)):
if pos[r] < len(lists[r]) and (best == -1 or lists[r][pos[r]] < lists[best][pos[best]]):
best = r
merged.append(lists[best][pos[best]])
pos[best] += 1
return mergedk 個の先頭要素からなる最小ヒープ
考え方
各回で変更されるのは先頭の要素だけですが、走査のたびにすべての k 個を読み直します。最小ヒープを使うと、最小の要素が先頭になります。各要素にその行を示すタグを付けておけば、(value, row) を取り出した後、どの行の次の値を追加すればよいかがわかります。ヒープには、値が残っている各行につき常に1つの要素が入るので、先頭が未使用の値のうち最小のものになります。
最初の例では、ヒープは2、1、3から始まります。1を取り出して4を追加します。2を取り出して6を追加します。3を取り出して5を追加します。4を取り出して10を追加します。5を取り出します。3行目は値を使い切っているので、何も追加しません。6を取り出して9を追加し、次に9と10を取り出します。
各値は1回ずつ追加され、1回ずつ取り出されます。また、ヒープに入る要素は最大でも k 個です。つまり、O(log k) の操作を 2N 回行います。N = k = 10^4 の場合、ステップ数は 3 × 10^5 未満です。ヒープに必要なメモリは O(k) です。
アルゴリズム
- 各行
rについて、(lists[r][0], r)を値で順序付けされた最小ヒープに追加します。 - ヒープが空でない間、最小のペア
(value, r)を取り出し、valueを結果に追加します。 - 行
rに次の値がある場合、その値とrを一緒に追加します。 - ヒープが空になったら、結果を返します。
import heapq
def mergeKLists(lists):
# The heap holds one (value, row) pair per row that still has values: that row's head.
heap = [(row[0], r) for r, row in enumerate(lists)]
heapq.heapify(heap)
nxt = [1] * len(lists) # nxt[r]: index of row r's next value, not yet in the heap
merged = []
while heap:
value, r = heapq.heappop(heap) # the smallest head of all rows
merged.append(value)
if nxt[r] < len(lists[r]):
heapq.heappush(heap, (lists[r][nxt[r]], r)) # row r's new head takes its place
nxt[r] += 1
return merged
落とし穴と境界ケース
- 値だけをそのままプッシュする。ポップした後、どの行を進めればよいか分からなくなる。値と一緒に行も格納する。
- 誤って最大ヒープを使う。C++ の
priority_queueと Rust のBinaryHeapは最大値を先頭に置くため、greater<>またはReverseを使う。 - 最初にすべての値をプッシュする。結果はソートされたままだが、ヒープのサイズが
N件まで増え、処理量がO(N log N)になる。 - 短い行の末尾を越えて読み取る。最初の例では3行目の値は2つなので、次の値をプッシュする前に確認する。
- 重複を取り除く。2つ目の例にある3つの5はそれぞれ別の値であり、すべて結果に含める。
よくある質問3
Merge k Sorted Lists の時間計算量はどれくらいですか？
最小ヒープを使うと、計算量は O(N log k) です。ここで、N は値の総数、k はリストの数です。各値は一度ずつプッシュおよびポップされ、ヒープに保持される要素数は最大でも k です。出力以外に必要な追加メモリは O(k) です。
すべての値をまとめて並べ替えてみてはどうでしょうか？
それは正しく、時間計算量は O(N log N) です。リストがすでにソートされていることを考慮していないため、ヒープでは log k で済むところ、値ごとに log N のコストがかかります。また、すべての値を一度にメモリ上に保持する必要があります。ヒープなら、ストリームとして届くリストもマージできます。
ヒープを使わずに、k 個のソート済みリストをマージできますか？
はい、分割統治法を使います。2つのリストをマージする方法でリストをペアごとにマージし、次にその結果をペアごとにマージしていきます。これを繰り返すと、各回ですべての値に1回ずつ触れる log k 回の処理となるため、O(N log k) です。リストを1つずつ大きくなる結果にマージしていく方法は、最初の方の値がマージのたびに再コピーされるため、より遅く、O(N·k) です。
似た問題
同じ考え方を使う問題です。2〜3問解くとパターンが身につきます。
def mergeKLists(lists):
# ここにコードを書いてくださいテスト #1
テスト #2
テスト #3
引数
lists = [[2, 6, 9], [1, 4, 10], [3, 5]]
期待される出力
[1, 2, 3, 4, 5, 6, 9, 10]