[路飞]_leetcode_373.查找和最小的K对数字

给定两个以升序排列的整数数组 nums1 和 nums2 , 以及一个整数 k 。

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

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

解题思路

这一题我们可以通过大顶堆维护二维数组中的数据大小值。

1. 首先先将二维数组 [nums1[i], nums2[i]] 插入到小顶堆中
2. 小顶堆通过判断数组内值的和来计算大小。

nums1: [1,7,11]
nums2: [2,4,6]
k = 3

[1, 2] 3
            [1, 2]

[1, 4] 5
            [1, 4]
           /
       [1, 2]
[1, 6] 7
            [1, 6]
           /     \
       [1, 2]    [1, 4]

[7, 2] 9
            [7, 2]
           /     \
       [1, 6]    [1, 4]
        /
    [1, 2]

因为堆的长度大于 k， 需要维护最小的k个数，所以需要将堆中最大值pop出去
            [1, 6]
           /     \
       [1, 2]    [1, 4]

[7, 11] 18
            [7, 11]
           /     \
       [1, 6]    [1, 4]
        /
    [1, 2]
和上面的一样操作
...

代码

class Heap {
    constructor() {
        this.data = [];
    }

    swap(i, j) {
        [this.data[i], this.data[j]] = [this.data[j], this.data[i]]
    }

    compartor(i, j) {
        return (this.data[i][0] + this.data[i][1]) - (this.data[j][0] + this.data[j][1]);
    }

    size() {
        return this.data.length;
    }

    peek() {
        if (!this.size()) return null;
        return this.data[0]
    }

    offer(val) {
        this.data.push(val);
        this.bubbleUp(this.data.length-1);
    }

    pop() {
        if (!this.size()) return null;
        let res = this.data[0];
        if (this.size() > 1) {
            this.data[0] = this.data.pop();
            this.bubbleDown(0);
        } else {
            this.data.pop();
        }
        return res;
    }

    bubbleUp(index) {
        while(index) {
            let parent = (index - 1) >> 1;
            if (this.compartor(index, parent) > 0) {
                this.swap(index, parent);
                index = parent;
            } else {
                break;
            }
        }
    }

    bubbleDown(index) {
        let last = this.data.length - 1;
        while(index < last) {
            let temp = index, left = index * 2 + 1, right = index * 2 + 2;
            if (left <= last && this.compartor(left, temp) > 0) {
                temp = left;
            }
            if (right <= last && this.compartor(right, temp) > 0) {
                temp = right;
            }
            if (temp == index) break;
            this.swap(temp, index);
            index = temp;
        }
    }
}
var kSmallestPairs = function(nums1, nums2, k) {
    let h = new Heap();
    for (x of nums1) {
        for (y of nums2) {
            let temp = [];
            temp[0] = x; temp[1] = y;
            h.offer(temp);
            if (h.size() > k) h.pop(); 
        }
    }
    return h.data;
};