[路飞]_373. 查找和最小的K对数字_堆

140 阅读2分钟

题目地址
堆-二叉堆数据结构

题目描述

给定两个以升序排列的整数数组 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]

示例 3:

输入: nums1 = [1,2], nums2 = [3], k = 3 
输出: [1,3],[2,3]
解释: 也可能序列中所有的数对都被返回:[1,3],[2,3]

解题思路

  • 求最值问题,一般都是用堆,因为本题是求最小值所以需要维护一个大顶堆
  • 双层for循环两个数组,判断1,如果堆中的元素长度小于k,那么就继续push,并维护一个shiftUp大顶堆上浮。判断2,如果sum1[i] + sum2[j]中的元素小于等于堆顶元素,那么弹出堆顶元素,并把sum1[i] + sum2[j]加入堆中,并维护一个大顶堆下沉。
代码:
var kSmallestPairs = function(nums1, nums2, k) {
    let heap = [];
    for(let i = 0; i < nums1.length; i++){
        for(let j = 0; j < nums2.length; j++){
            if(heap.length < k){
                heap.push([nums1[i], nums2[j]]);
                shiftUp(heap, heap.length - 1);
            }else if(sum(heap[0]) >= (nums1[i] + nums2[j])){
                heap[0] = [nums1[i], nums2[j]];
                shiftDown(heap, 0)
            }
        }
    }
    return heap;
};

function sum(arr){
    return arr[0] + arr[1];
}

function swap(heap, index, parent){
    [heap[index], heap[parent]] = [heap[parent], heap[index]];
}

function shiftUp(heap, index){
    let parent = (index - 1) / 2 | 0;
    if(sum(heap[index] > sum(heap[parent]))){
        swap(heap, index, parent);
        shiftUp(heap, parent);
    }
}

function shiftDown(heap, index){
    let left = index * 2 + 1;
    if(left >= heap.length) return;
    if(left + 1 < heap.length && sum(heap[left]) < sum(heap[left + 1])){
        left = left + 1;
    }
    if(sum(heap[index] <= sum(heap[left]))){
        swap(heap, index, left);
        shiftDown(heap, left);
    }
}
代码:
var kSmallestPairs = function(nums1, nums2, k) { 
    let n = nums1.length; 
    let m = nums2.length; 
    k = k > n * m ? n * m : k; // tails[i]表示nums2中与nums[i]相加的最小值下标 
    let tails = []; 
    for(let i = 0; i < n; i++){ 
        tails[i] = 0; 
    } 
    let result = [];
    while(k > 0){ 
        let minIndex = 0; 
        let minCount = 2147483647; 
        for(let i =0; i< tails.length; i++){ 
            if(tails[i] >= m){ 
                continue; 
            } 
            if(nums1[i] + nums2[tails[i]] < minCount){ 
                minIndex = i; 
                minCount = nums1[i] + nums2[tails[i]]; 
            } 
        } 
        result.push([nums1[minIndex], nums2[tails[minIndex]]]); 
        tails[minIndex]++; 
        k--; 
    } 
    return result; 
};