Medium

题目描述

给定两个以 非递减顺序排列 的整数数组 nums1 和 nums2 ，以及一个整数 k。

定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2。

请找到和最小的 k 个数对 (u1,v1), (u2,v2) … (uk,vk)。

示例 1：

输入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
输出: [[1,2],[1,4],[1,6]]
解释: 返回序列中的前 3 对数字:
     [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

示例 2：

输入: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
输出: [1,1],[1,1]
解释: 返回序列中的前 2 对数字:
     [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

提示：

  • 1 <= nums1.length, nums2.length <= 10^5
  • -10^9 <= nums1[i], nums2[i] <= 10^9
  • nums1 和 nums2 均为升序排列
  • 1 <= k <= 10^4
  • k <= nums1.length * nums2.length

解题思路

本题需要找到和最小的 k 对数字。由于两个数组都是有序的，我们可以使用**堆（优先队列）**来高效解决。

核心思路：

  1. 使用最小堆存储候选数对，按照和的大小排序
  2. 初始时将 (nums1[0], nums2[0]) 到 (nums1[min(k-1, len1-1)], nums2[0]) 加入堆
  3. 每次从堆中取出和最小的数对作为结果
  4. 取出 (i, j) 后，如果 j+1 < len2，则将 (i, j+1) 加入堆

为什么这样做？

  • 由于数组有序，对于任意 i，最小的数对一定是 (nums1[i], nums2[0])
  • 当我们取出 (i, j) 后，下一个可能的最小数对是 (i, j+1)
  • 使用堆保证每次都能取到当前候选中和最小的数对

时间复杂度优化： 初始时不需要将所有 nums1 的元素都加入堆，只需要前 k 个（或全部，如果 len1 < k），这样可以避免不必要的计算。

代码实现

class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<vector<int>> result;
        if (nums1.empty() || nums2.empty()) return result;
        
        // 最小堆：存储 {和, {i, j}}
        priority_queue<pair<int, pair<int, int>>, vector<pair<int, pair<int, int>>>, greater<pair<int, pair<int, int>>>> pq;
        
        // 初始化：将前min(k, nums1.size())个元素与nums2[0]组成的数对加入堆
        for (int i = 0; i < min(k, (int)nums1.size()); i++) {
            pq.push({nums1[i] + nums2[0], {i, 0}});
        }
        
        while (k-- > 0 && !pq.empty()) {
            auto top = pq.top();
            pq.pop();
            
            int i = top.second.first;
            int j = top.second.second;
            
            result.push_back({nums1[i], nums2[j]});
            
            // 如果j+1还在范围内，将(i, j+1)加入堆
            if (j + 1 < nums2.size()) {
                pq.push({nums1[i] + nums2[j + 1], {i, j + 1}});
            }
        }
        
        return result;
    }
};

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        if not nums1 or not nums2:
            return []
        
        import heapq
        
        # 最小堆：存储 (和, i, j)
        heap = []
        
        # 初始化：将前min(k, len(nums1))个元素与nums2[0]组成的数对加入堆
        for i in range(min(k, len(nums1))):
            heapq.heappush(heap, (nums1[i] + nums2[0], i, 0))
        
        result = []
        
        while k > 0 and heap:
            sum_val, i, j = heapq.heappop(heap)
            result.append([nums1[i], nums2[j]])
            
            # 如果j+1还在范围内，将(i, j+1)加入堆
            if j + 1 < len(nums2):
                heapq.heappush(heap, (nums1[i] + nums2[j + 1], i, j + 1))
            
            k -= 1
        
        return result

public class Solution {
    public IList<IList<int>> KSmallestPairs(int[] nums1, int[] nums2, int k) {
        var result = new List<IList<int>>();
        if (nums1.Length == 0 || nums2.Length == 0) return result;
        
        // 最小堆：存储 (和, i, j)
        var pq = new PriorityQueue<(int sum, int i, int j), int>();
        
        // 初始化：将前min(k, nums1.Length)个元素与nums2[0]组成的数对加入堆
        for (int i = 0; i < Math.Min(k, nums1.Length); i++) {
            pq.Enqueue((nums1[i] + nums2[0], i, 0), nums1[i] + nums2[0]);
        }
        
        while (k-- > 0 && pq.Count > 0) {
            var (sum, i, j) = pq.Dequeue();
            result.Add(new List<int> { nums1[i], nums2[j] });
            
            // 如果j+1还在范围内，将(i, j+1)加入堆
            if (j + 1 < nums2.Length) {
                pq.Enqueue((nums1[i] + nums2[j + 1], i, j + 1), nums1[i] + nums2[j + 1]);
            }
        }
        
        return result;
    }
}

var kSmallestPairs = function(nums1, nums2, k) {
    const result = [];
    const heap = [];
    
    // Add first k pairs from nums1[0] with nums2[0...k-1]
    for (let j = 0; j < Math.min(k, nums2.length); j++) {
        heap.push([nums1[0] + nums2[j], 0, j]);
    }
    
    // Min heap implementation
    const parent = i => Math.floor((i - 1) / 2);
    const left = i => 2 * i + 1;
    const right = i => 2 * i + 2;
    
    const heapifyUp = (i) => {
        while (i > 0 && heap[i][0] < heap[parent(i)][0]) {
            [heap[i], heap[parent(i)]] = [heap[parent(i)], heap[i]];
            i = parent(i);
        }
    };
    
    const heapifyDown = (i) => {
        let minIdx = i;
        const l = left(i);
        const r = right(i);
        
        if (l < heap.length && heap[l][0] < heap[minIdx][0]) minIdx = l;
        if (r < heap.length && heap[r][0] < heap[minIdx][0]) minIdx = r;
        
        if (minIdx !== i) {
            [heap[i], heap[minIdx]] = [heap[minIdx], heap[i]];
            heapifyDown(minIdx);
        }
    };
    
    const push = (item) => {
        heap.push(item);
        heapifyUp(heap.length - 1);
    };
    
    const pop = () => {
        if (heap.length === 0) return null;
        const min = heap[0];
        heap[0] = heap[heap.length - 1];
        heap.pop();
        if (heap.length > 0) heapifyDown(0);
        return min;
    };
    
    // Build initial heap
    for (let i = heap.length - 1; i >= 0; i--) {
        heapifyDown(i);
    }
    
    while (result.length < k && heap.length > 0) {
        const [sum, i, j] = pop();
        result.push([nums1[i], nums2[j]]);
        
        // Add next pair from same nums1 index if exists
        if (i + 1 < nums1.length) {
            push([nums1[i + 1] + nums2[j], i + 1, j]);
        }
    }
    
    return result;
};
在 LeetCode 上运行： 力扣中文版 LeetCode 国际版

复杂度分析

复杂度类型复杂度说明
时间复杂度O(k log k)最多进行k次堆操作，每次操作O(log k)
空间复杂度O(k)堆的大小最多为k

相关题目