LeetCode LCR 061. 查找和最小的 K 对数字

题目描述

LCR 061. 查找和最小的 K 对数字

思路分析

解法一:最小堆(推荐)

核心思路

  • (i, 0) 入堆,表示 nums1[i] + nums2[0]
  • 每次弹出当前最小和的 pair,并将 (i, j+1) 入堆。
  • 重复 k 次即可得到前 k 小的数对。

复杂度分析

  • 时间复杂度:O(k log n),其中 n 表示 nums1 长度。
  • 空间复杂度:O(n),堆中最多保存 nums1 个元素。
class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> res = new ArrayList<>();
        if (nums1.length == 0 || nums2.length == 0 || k == 0) {
            return res;
        }

        PriorityQueue<int[]> pq = new PriorityQueue<>(
                (a, b) -> (nums1[a[0]] + nums2[a[1]]) - (nums1[b[0]] + nums2[b[1]])
        );

        int n1 = nums1.length;
        int n2 = nums2.length;
        for (int i = 0; i < Math.min(k, n1); i++) {
            pq.offer(new int[]{i, 0});
        }

        // 每次弹出最小和并扩展下一列
        while (!pq.isEmpty() && k-- > 0) {
            int[] cur = pq.poll();
            int i = cur[0];
            int j = cur[1];

            res.add(Arrays.asList(nums1[i], nums2[j]));
            if (j + 1 < n2) {
                pq.offer(new int[]{i, j + 1});
            }
        }

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

	h := &PairHeap{nums1: nums1, nums2: nums2}
	heap.Init(h)
	for i := 0; i < len(nums1) && i < k; i++ {
		heap.Push(h, Pair{i: i, j: 0})
	}

	// 每次弹出最小和并扩展下一列
	for h.Len() > 0 && k > 0 {
		cur := heap.Pop(h).(Pair)
		res = append(res, []int{nums1[cur.i], nums2[cur.j]})
		if cur.j+1 < len(nums2) {
			heap.Push(h, Pair{i: cur.i, j: cur.j + 1})
		}
		k--
	}

	return res
}

type Pair struct {
	i int
	j int
}

type PairHeap struct {
	data  []Pair
	nums1 []int
	nums2 []int
}

func (h PairHeap) Len() int { return len(h.data) }
func (h PairHeap) Less(i, j int) bool {
	a := h.data[i]
	b := h.data[j]
	sumA := h.nums1[a.i] + h.nums2[a.j]
	sumB := h.nums1[b.i] + h.nums2[b.j]
	return sumA < sumB
}
func (h PairHeap) Swap(i, j int) { h.data[i], h.data[j] = h.data[j], h.data[i] }
func (h *PairHeap) Push(x any)   { h.data = append(h.data, x.(Pair)) }
func (h *PairHeap) Pop() any {
	old := h.data
	n := len(old)
	val := old[n-1]
	h.data = old[:n-1]
	return val
}

相似题目

题目 难度 考察点
378. 有序矩阵中第 K 小的元素 中等
719. 找出第 K 小的数对距离 困难 二分 + 双指针
347. 前 K 个高频元素 中等
373. 查找和最小的 K 对数字 中等
786. 第 K 个最小的质数分数 中等
本文作者:
本文链接: https://hgnulb.github.io/blog/2026/44061981
版权声明: 本博客所有文章除特别声明外,均采用 CC BY-NC-SA 4.0 许可协议,转载请注明出处!