题目描述

✅ 373. 查找和最小的 K 对数字

image-20260928235922541

image-20260928235922542

题意分析

从两个非递减数组中各取一个元素组成数对,找出数对和最小的前 k 对,返回每对中的两个数值。一次选择由两个数组下标共同确定,同一个元素可以与另一数组中的多个元素配对,并不是每个元素只能使用一次。

数组可能含重复值。来自不同下标组合的数对即使数值相同,仍然是不同的选择,输出中要保留相应出现次数。和相等的候选不要求固定先后次序,只要选出的前 k 对没有漏掉更小的和即可。

解法:小根堆多路归并

核心思路

[!blue]

直接生成所有数对需要处理 m × n 种组合,即使只需要很少的结果也会浪费大量工作。利用数组有序的条件,固定 nums1 的下标 i,让第二个下标依次取 0、1、2…,这一行数对的和就按非递减顺序排列。问题因此变成从多条有序序列中取出最小的前 k 项。

每一行只需要把尚未输出的第一项放进小根堆,状态保存数对和以及下标 (i, j)。行内其他项都不小于这一项,暂时无需入堆;所有行的首个候选中最小的堆顶,也就不大于任何尚未输出的数对,是全局下一项。

弹出 (i, j) 后,只把同一行的 (i, j + 1) 补入堆,继续维持“每条活动行一个最小未输出候选”的不变量。每个下标对只会通过前一列进入一次,所以不需要额外的访问集合;不同下标对应的相同值则应当照常输出。

初始化也不必开启全部 m 行,只需前 min(m, k) 行。如果某行下标不小于 k,排在它前面的 k 个行首已经都不大于它的行首,而这一行后续数对更不会更小。因此前 k 个结果总能从保留的行里选出;遇到相同的和时,选择较早行中的候选仍然合法。

得到 k 对后停止。代码也检查堆是否为空,使输入为空、k 为零或实际数对不足时能够自然结束,不会从空堆中取值。

解题步骤

  1. 任一数组为空或 k == 0 时返回空列表。
  2. 令 rows = min(nums1.length, k),把下标对 (i, 0) 的和与坐标加入小根堆,其中 i 遍历前 rows 行。
  3. 弹出和最小的状态 (i, j),把对应的两个数组值加入答案。
  4. 若这一行还有下一列,将 (i, j + 1) 加入堆;其他行的候选保持不变。
  5. 反复弹出并补入,直到已经输出 k 对或堆变为空。

代码实现

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> result = new ArrayList<>();

        if (nums1.length == 0 || nums2.length == 0 || k == 0) {
            return result;
        }

        PriorityQueue<long[]> heap = new PriorityQueue<>((a, b) -> Long.compare(a[0], b[0]));
        // 后续行的首项前已有足够候选,只开启前 k 行
        int rows = Math.min(nums1.length, k);

        for (int i = 0; i < rows; i++) {
            heap.offer(new long[] {
                (long) nums1[i] + nums2[0],
                i,
                0
            });
        }

        while (k-- > 0 && !heap.isEmpty()) {
            long[] state = heap.poll();
            int i = (int) state[1];
            int j = (int) state[2];

            result.add(Arrays.asList(nums1[i], nums2[j]));

            // 只推进刚弹出的那一行,维持每行一个候选
            if (j + 1 < nums2.length) {
                heap.offer(new long[] {
                    (long) nums1[i] + nums2[j + 1],
                    i,
                    j + 1
                });
            }
        }

        return result;
    }
}
import "container/heap"

type pair struct {
    sum int64
    i   int
    j   int
}

type pairHeap []pair

func (h pairHeap) Len() int { return len(h) }

func (h pairHeap) Less(i, j int) bool { return h[i].sum < h[j].sum }

func (h pairHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }

func (h *pairHeap) Push(x any) {
    *h = append(*h, x.(pair))
}

func (h *pairHeap) Pop() any {
    old := *h
    item := old[len(old)-1]
    *h = old[:len(old)-1]
    return item
}

func kSmallestPairs(nums1, nums2 []int, k int) [][]int {
    result := make([][]int, 0)
    if len(nums1) == 0 || len(nums2) == 0 || k == 0 {
        return result
    }

    // 后续行的首项前已有足够候选,只开启前 k 行
    rows := len(nums1)
    if k < rows {
        rows = k
    }
    h := &pairHeap{}
    heap.Init(h)
    for i := 0; i < rows; i++ {
        heap.Push(h, pair{
            sum: int64(nums1[i]) + int64(nums2[0]),
            i:   i,
        })
    }

    for k > 0 && h.Len() > 0 {
        current := heap.Pop(h).(pair)
        result = append(result, []int{
            nums1[current.i],
            nums2[current.j],
        })
        k--

        // 只推进刚弹出的那一行,维持每行一个候选
        if current.j+1 < len(nums2) {
            nextJ := current.j + 1
            heap.Push(h, pair{
                sum: int64(nums1[current.i]) + int64(nums2[nextJ]),
                i:   current.i,
                j:   nextJ,
            })
        }
    }
    return result
}

复杂度分析

  • 时间复杂度:O((r + p) log(r + 1)),其中 r = min(m, k) 是初始化行数,p 为实际输出数量。代码逐项入堆初始化,随后每次输出做一次弹出、至多一次补入,堆大小始终不超过 r。
  • 空间复杂度:辅助空间 O(r),用于堆中的活动行候选;返回的 p 对结果另占 O(p)。

关键点总结

[!green]

  • 固定一个数组下标即可得到一条有序数对序列,堆只需比较各行尚未输出的第一项。
  • 弹出哪一行就推进哪一行,保证覆盖完整且每个下标对只有一个生成来源。
  • 第二个数组有序保证行内归并,第一个数组有序保证只初始化前 k 行不会漏掉更小结果。
  • 排序依据是数对和,重复的值不代表重复的下标组合,不能按值去重。

易错点总结

[!yellow]

  • 每次同时生成右侧和下一行候选:同一下标对可能从两个方向进入堆;本解法统一从行首开始,只向右推进。
  • 推进其他行:弹出后必须补上同一行的下一列,否则会漏掉该行本应更早出现的候选。
  • 把相同数值数对去重:不同下标组合仍是不同结果,去重会减少合法出现次数。
  • 使用两数和相减作比较器:即使每个和能装入整数,两个和的差也可能溢出,代码使用宽整数与安全比较。
  • 只看结果数量,不检查堆空:所有可用数对耗尽后需要停止,不能继续取堆顶。

相似题目

题目 难度 关联与区别
23. 合并 K 个升序链表 困难 同样做多路有序归并,本题每个固定的第一数组元素形成一条有序数对和序列。
378. 有序矩阵中第 K 小的元素 中等 两数组数对和形成行列有序矩阵,可对比求前k个对象与只求第k小数值的区别。
转载与许可
作者
链接 https://hgnulb.github.io/blog/2026/54475263
许可 本博客所有文章除特别声明外,均采用 CC BY-NC-SA 4.0 许可协议,转载请注明出处!