10. Find K Pairs with Smallest Sums

読了目安約2分

Heap / LeetCode 373 (Medium) — 2 つのソート済み配列から、和が小さい順に K 個のペアを取り出す。

この章の目次

問題

LeetCode 373. Find K Pairs with Smallest Sums（Medium / カテゴリ: Heap）

昇順にソートされた 2 つの整数配列 nums1、nums2 と整数 k が与えられます。 nums1 から 1 つ、nums2 から 1 つ取ったペア [u, v] のうち、和 u + v が小さい方から k 個を返してください。ペアの総数が k 未満のときは全ペアを返します。

考え方

ヒント 1

全ペアを作ってソートすると O(mn log(mn)) で、配列が大きいと間に合いません。最小の和のペアは必ず (nums1[0], nums2[0]) です。取り出した結果から候補を少しずつ広げる方法を考えます。

ヒント 2

（和, i, j）を優先度付きキュー（最小ヒープ）で管理します。先に (i, 0) を i = 0..<min(k, m) の分だけ積んでおき、(i, j) を取り出すたびに (i, j + 1) だけを新たに積みます。これで訪問済みの管理なしに、和が小さい順へ k 回取り出せます。

Swift 実装のポイント

標準ライブラリにヒープはないので、(sum: Int, i: Int, j: Int) のタプルを持つ二分ヒープを実装します。比較は sum だけで十分です。
k がペア総数より大きい場合は、ヒープが空になった時点で打ち切ります。while result.count < k, let item = heap.pop() の形が簡潔です。
取り出した (i, j) から積み直すのは (i, j + 1) の 1 つだけです。j + 1 < nums2.count の範囲チェックを忘れないでください。

模範解答

Swift
struct MinHeap {
    private var items: [(sum: Int, i: Int, j: Int)] = []

    mutating func push(_ item: (sum: Int, i: Int, j: Int)) {
        items.append(item)
        var i = items.count - 1
        while i > 0 {
            let parent = (i - 1) / 2
            if items[parent].sum <= items[i].sum { break }
            items.swapAt(parent, i)
            i = parent
        }
    }

    mutating func pop() -> (sum: Int, i: Int, j: Int)? {
        guard let first = items.first else { return nil }
        items[0] = items[items.count - 1]
        items.removeLast()
        var i = 0
        while true {
            let left = 2 * i + 1
            let right = 2 * i + 2
            var smallest = i
            if left < items.count && items[left].sum < items[smallest].sum { smallest = left }
            if right < items.count && items[right].sum < items[smallest].sum { smallest = right }
            if smallest == i { break }
            items.swapAt(i, smallest)
            i = smallest
        }
        return first
    }
}

class Solution {
    func kSmallestPairs(_ nums1: [Int], _ nums2: [Int], _ k: Int) -> [[Int]] {
        if nums1.isEmpty || nums2.isEmpty || k <= 0 {
            return []
        }
        var heap = MinHeap()
        for i in 0..<min(k, nums1.count) {
            heap.push((nums1[i] + nums2[0], i, 0))
        }
        var result: [[Int]] = []
        while result.count < k, let item = heap.pop() {
            result.append([nums1[item.i], nums2[item.j]])
            if item.j + 1 < nums2.count {
                heap.push((nums1[item.i] + nums2[item.j + 1], item.i, item.j + 1))
            }
        }
        return result
    }
}

計算量: O(k log k) 時間、O(k) 追加メモリ。