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

记录 1 道算法题

查找和最小的K对数字

leetcode-cn.com/problems/fi…

上来先整个暴力解法。但是不能太暴力，不然会超时。所以用了一个小的规律。如果是两个数组个数都是大于k个，那么我们最多取到 k * k 个就可以了。哪怕极端一点，比如 [1, 5, 6], [2, 3, 4] 取 (1,2) (1,3) (1, 4)。也是不会超过边长 k 的。

    function kSmallestPairs(nums1, nums2, k) {
        let n1 = nums1.length > k ? k : nums1.length
        let n2 = nums2.length > k ? k : nums2.length
        let result = []
        for(let i = 0; i < n1; i++) {
            for(let j = 0; j < n2; j++) {
                result.push([nums1[i], nums2[j]])
            }
        }

        result.sort((a, b) => a[0] + a[1] - b[0] - b[1])

        return result.slice(0, k)
    }

上面根据数组的长度不同，最多维护了 k * k 个，然后进行排序。如果使用堆来进行计算的话，就只需要维护 k 个就可以了。最小的 k 对数字，用最大堆。

另外一篇文章详细的介绍了堆的优化计算

不过这里用的是最大堆，所以 this.compare 要改成降序 (a, b) => b - a

然后堆的冒泡和下沉道理都是相同的，使用搜索二叉树。

再加一个小优化就是，因为堆顶是这 k 对元素里面最大的，所以如果比堆顶的数还要大，就可以直接排除。

完整代码如下：

    function kSmallestPairs(nums1, nums2, k) {
        const heap = new Heap()
        // 依然最多是 k * k 个
        let n1 = nums1.length > k ? k : nums1.length
        let n2 = nums2.length > k ? k : nums2.length
        for (let i = 0; i < n1; i++) {
          for (let j = 0; j < n2; j++) {
            const val = nums1[i] + nums2[j]
            // 如果堆以及满了，就和堆顶的数比较，比堆顶的数大就跳过
            if(heap.size() === k && heap.data[0].val < val ) {
                continue
            }
            heap.push({ data: [nums1[i], nums2[j]], val })
            // 当超出了就要踢出最大的，进行重新下沉
            if (heap.size() > k) {
              heap.pop()
            }
          }
        }

        return heap.data.reduce((a, b) => {
          a.push(b.data)
          return a
        }, [])
      }
      
      class Heap {
        constructor() {
          this.data = []
          this.compare = (a, b) => b - a
        }

        size() {
          return this.data.length
        }

        swap(n1, n2) {
          const { data } = this
          const temp = data[n1]
          data[n1] = data[n2]
          data[n2] = temp
        }

        push(val) {
          this.data.push(val)
          this.bubblingUp(this.size() - 1)
        }

        pop() {
          if (this.size() === 0) return null
          const { data } = this
          const discard = data[0]
          const newMember = data.pop()
          if (this.size() > 0) {
            data[0] = newMember
            this.bubblingDown(0)
          }

          return discard
        }

        bubblingUp(index) {
          while (index > 0) {
            const parent = (index - 1) >> 1
            const { data } = this
            if (this.compare(data[index].val, data[parent].val) < 0) {
              this.swap(parent, index)
              index = parent
            } else {
              break
            }
          }
        }

        bubblingDown(index) {
          const { data } = this
          const last = this.size() - 1
          while (true) {
            const left = index * 2 + 1
            const right = index * 2 + 2
            let parent = index
            if (left <= last && this.compare(data[left].val, data[parent].val) < 0) {
              parent = left
            }
            if (right <= last && this.compare(data[right].val, data[parent].val) < 0) {
              parent = right
            }
            if (index !== parent) {
              this.swap(index, parent)
              index = parent
            } else {
              break
            }
          }
       }
    }

结束