[路飞]和最小的K对数

今天我们来做 剑指 Offer II 061. 和最小的 k 个数对

题目

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

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

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

图示

解析

• 图示分析， 需要创建一个大顶堆，保持最多有K个元素
• 当新来的一对数的和值 小于我们大顶堆的根节点和值的时候，再去push
• 如果堆大小 大于K的时候 把最大的根顶元素pop出
• 堆元素比较大小： 根据元素和值的大小

代码

var kSmallestPairs = function(nums1, nums2, k) {
  // 创建堆
  const heap = new Heap()

  for (let i = 0; i < nums1.length; i++) {
    for (let j = 0; j < nums2.length; j++) {
      // 进堆的条件
      if (heap.size() < k || defaultCompare(heap.top(), [nums1[i], nums2[j]])) {
        heap.push([nums1[i], nums2[j]])
        if (heap.size() > k) {
          heap.pop()
        }
      } else {
        // 升序 所以可以直接把第二循环break掉
        break
      }
    }
  }

  return heap.data
};

// 定义比较方法
function defaultCompare(a, b) {
  return a[0] + a[1] > b[0] + b[1]
}

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

  size() {
    return this.data.length
  }

  top() {
    return this.data[0]
  }

  push(val) {
    this.data.push(val)

    let index = this.size() - 1

    while (index > 0) {
      const parentIndex = Math.floor((index - 1) / 2)
      const condition = this.compare(this.data[index], this.data[parentIndex])
      if (condition) {
        this.swap(index, parentIndex)
        index = parentIndex
      } else {
        break
      }
    }

    return this.top()
  }

  pop() {
    if (this.size() === 0) return
    const lastOne = this.data.pop()
    let r
    if (this.size() > 0) {
      r = this.data[0]
      this.data[0] = lastOne

      let index = 0

      while (index * 2 + 1 < this.size()) {
        const leftIndex = index * 2 + 1
        const rightIndex = index * 2 + 2

        let findIndex = index

        const conditionL = this.compare(
          this.data[leftIndex],
          this.data[findIndex]
        )
        if (conditionL) {
          findIndex = leftIndex
        }

        const conditionR =
          rightIndex < this.size() &&
          this.compare(this.data[rightIndex], this.data[findIndex])

        if (conditionR) {
          findIndex = rightIndex
        }

        if (findIndex > index) {
          this.swap(findIndex, index)
          index = findIndex
        } else {
          break
        }
      }
    }
    return r === undefined ? lastOne : r
  }

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

优化

• 因为升序 在一层循环中： 可以看出 如果 nums1[i] + nums2[0] > heap.top()[0] + heap.top()[1]的时候也可以break掉后面

...
// 在第一层循环开始坐下判断
if (heap.size() >= k && defaultCompare([nums1[i], nums2[0]], heap.top())) break
...

• 简单题重拳出击