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

130 阅读1分钟

题目介绍

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

提示:

1 <= nums1.length, nums2.length <= 105 -109 <= nums1[i], nums2[i] <= 109 nums1 和 nums2 均为升序排列 1 <= k <= 1000

来源:力扣(LeetCode) 链接:leetcode-cn.com/problems/fi… 著作权归领扣网络所有。商业转载请联系官方授权,非商业转载请注明出处。

解题思路

  1. 求最值问题,也是用堆的思路解决,实现一个最大堆
  2. 遍历nums1和nums2,如果对的size小于k 或者 当前nums1的元素和当前nums2的元素相加小于堆顶值时候才压入堆中(这是个优化点,不然会导致超时)
  3. 如果堆的size大于k值,弹出堆顶元素

代码

/**
 * @param {number[]} nums1
 * @param {number[]} nums2
 * @param {number} k
 * @return {number[][]}
 */
class Heap {
  constructor() {
    this.data = []
  }
  // 向下调
  shiftDown(ind = 0) {
    let n = this.size() - 1
    while (ind * 2 + 1 <= n) {
      let tmp = ind
      if (this.getVal(this.data[tmp]) < this.getVal(this.data[ind * 2 + 1])) tmp = ind * 2 + 1
      if (ind * 2 + 2 <= n && this.getVal(this.data[tmp]) < this.getVal(this.data[ind * 2 + 2]))
        tmp = ind * 2 + 2
      if (tmp === ind) break
      this.data[tmp] = [this.data[ind], (this.data[ind] = this.data[tmp])][0]
      ind = tmp
    }
  }
  getVal(arr) {
    return arr[0] + arr[1]
  }
  // 向上调整
  shiftUp(idx) {
    let pIdx = null
    while (
      ((pIdx = Math.floor((idx - 1) / 2)),
      idx && this.getVal(this.data[pIdx]) < this.getVal(this.data[idx]))
    ) {
      this.data[pIdx] = [this.data[idx], (this.data[idx] = this.data[pIdx])][0]
      idx = pIdx
    }
  }
  push(val) {
    this.data.push(val)
    this.shiftUp(this.size() - 1)
  }
  pop() {
    if (this.size() === 0) return
    if (this.size() === 1) {
      return this.data.pop()
    }
    this.data[0] = this.data.pop()
    this.shiftDown(0)
  }
  top() {
    return this.data[0]
  }
  size() {
    return this.data.length
  }
}
var kSmallestPairs = function(nums1, nums2, k) {
    const h = new Heap()
    for (let val1 of nums1) {
        for (let val2 of nums2) {
            // 只有比堆顶小的元素才可以放进去
            if (h.size() < k || (val2 + val1 < h.top()[0] + h.top()[1])) {
                h.push([val1, val2])
                if (h.size() > k) {
                    h.pop()
                }
            } else {
                break
            }
            
        }
    }
    return h.data
};