一日一练:最小的K对数字

134 阅读1分钟

给定两个以 升序排列 的整数数组 nums1 和 nums2 , 以及一个整数 k 。定义一对值 (u,v),其中第一个元素来自 nums1,第二个元素来自 nums2 。请找到和最小的 k 个数对 (u1,v1),  (u2,v2)  ...  (uk,vk) 。

虽然这里的两个数组都是升序排列,但是两个数组的数对和的大小是不确定的,所以用暴力扫描是无法实现的。最小或者最大的K个值一般都用最小堆或者最大堆。

这里可以用最小堆。需要注意的:

  1. 堆内不存数对和的大小,而是存nums1nums2的索引
  2. compare比较的是nums1nums2数对和的大小
class MinBinaryHeap {
  constructor(compare) {
    this.s = []
    this.len = 0
    this.compare = compare
  }

  push(el) {
    this.s[this.len] = el
    ++this.len
    if (this.len > 1) {
    this.swim(this.len - 1)
    }
  }

  pop() {
    if (this.len === 0) {
      return null
    }
    const top = this.s[0]

    if (this.len > 1) {
      this.s[0] = this.s[this.len - 1]
      this.sink(0)
    }
    this.len--
    return top
  }

  peek() {
    if (this.len === 0) {
      return null
    } else {
      return this.s[0]
    }
  }

  sink(idx) {
    const lastIdx = this.len - 1
    while (idx <= lastIdx) {
      const rightIdx = (idx + 1) * 2
      const leftIdx = rightIdx - 1
      if (leftIdx <= lastIdx) {
        let smallValueIdx = leftIdx
        if (rightIdx <= lastIdx && this.compare(this.s[smallValueIdx], this.s[rightIdx]) > 0) {
          smallValueIdx = rightIdx
        }
        if (this.compare(this.s[idx], this.s[smallValueIdx]) > 0) {
          this.swap(idx, smallValueIdx)
          idx = smallValueIdx
        } else {
          break
        }
      } else {
        break
      }
    }
  }

  swim(idx) {
    while (idx >= 0) {
      let parent = (idx - 1) >> 1
      if (parent >= 0 && this.compare(this.s[parent], this.s[idx]) > 0) {
        this.swap(parent, idx)
        idx = parent
      } else {
        break
      }
    }
  }

  swap(a, b) {
    let temp = this.s[a]
    this.s[a] = this.s[b]
    this.s[b] = temp
  }

  compare(a, b) {
    return a - b
  }
}

const kSmallestPairs = function(nums1, nums2, k) {
    const newPairsHeap = new MinBinaryHeap((a, b) => nums1[a[0]] + nums2[a[1]] - nums1[b[0]] - nums2[b[1]])
    const n1l = nums1.length
    const n2l = nums2.length
    // 保存k个值
    // 可以确定的是 最小的是[0, 0]
    for(let i = 0; i < Math.min(k, n1l); i++) {
        newPairsHeap.push([i, 0])
    }
    const ans = []
    while(newPairsHeap.len) {
        const top = newPairsHeap.pop()
        ans.push([nums1[top[0]], nums2[top[1]]])
        if (ans.length === k) {
            return ans
        }
        // 索引[0, 0]对应元素被加入到答案之后,将[0, 1]加入堆中,
        // [1, 0]之前已经加入了。因为是排过序的,所以最小值肯定是[0, 1]或者[1, 0]。
        // 这样索引较小的数对已加入堆中,保证下次为最小值
        if ((top[1] + 1) < n2l) {
            newPairsHeap.push([top[0], top[1] + 1])
        }
    }
    return ans
};