一日一练: 小顶堆

153 阅读3分钟

输入整数数组 arr ,找出其中最小的 k 个数。例如,输入4、5、1、6、2、7、3、8这8个数字,则最小的4个数字是1、2、3、4。

这道题的的第一种思路是sort + slice

这道题还可以使用最小堆或者说小顶堆实现。在js中并没有提供类似于Java中的工具类PriorityQueue的工具方法,这里需要自己去模拟。

堆是一种特别的完全二叉树。二叉树可以用链表实现,也可以用数组来实现。小顶堆一般用数组来模拟。

数组[1, 2, 3, 4, 5, 6, 7]

链表实现的时候,一般用node.leftnode.right来表示一个节点的左右子节点;而这里是用索引来表明他们的关系,索引0为根节点,索引1索引2可以看成是索引0左右子节点,他们之间的关系可以表示为:

  • 左子节点索引 = (父节点索引 + 1) * 2 - 1
  • 右子节点索引 = (父节点索引 + 1) * 2
  • 父节点索引 = 子节点索引或者右节点索引 >> 1

如下图所示:

image.png

小顶堆特性

每个结点的值都小于或等于其左右子结点的值,叫小顶堆。

实现思路

为了维持小顶堆的特性,需要在数组加入成员或者弹出元素成员之后对数组进行相对应的操作。

  • 加入元素:方式为push到队尾。将此时队尾元素与其父节点比较,如果比父节点小,就与父节点交互位置;之后继续与父节点比较。。。直到比父节点大或者已经到堆顶(索引为0)

  • 弹出元素:为了最小化修改,这里不直接shift。而是先记录堆顶值,然后将最后一位的值赋值给堆顶,最后针对堆顶做下沉比较。具体操作为:如果堆顶比其左右节点中小的节点更小,交互他们的位置;然后继续找左右子节点值小的节点交互位置。。。直到超过堆最大长度或者比左右节点都小。

    [1, 2, 3, 4, 5, 6, 7]为例:如果弹出堆顶元素1,将7的值赋值给索引0;此时堆顶为7,将7与左右节点2,3比较,2是较小的那个值,7再与2交互位置,此时堆顶已为最小的2;再让7与左右节点45比较,4为较小的那个值,交互位置,到此7没有节点,结束。重新调整之后为[2, 4, 3, 7, 5, 6]

代码

  • push: 入堆
// 手写堆
let s = []
// 这里不用s.length,可以在pop的时候不需要pop元素,只需要将长度-1
let len = 0 

function push(el) {
  s[len] = el
  if (++len > 1) {
    // 上浮刚加入的元素,保持最小堆 
    swim(len - 1)
  }
}
  • 💁 swim 入堆时的调度函数,保持最小堆的特性
function swim(idx) {
  while(idx >= 0) {
    let parent = (idx - 1) >> 1
    if (compare(parent, idx) > 0) {
      swap(parent, idx)
      idx = parent
    } else {
      break
    }
  }
}
  • pop 出堆,每次都是获取队内的最小值
function pop() {
  if (len === 0) {
    return null
  }
  // 堆顶
  const top = s[0]
  // 如果还有除堆顶元素以外的元素,将最后一个元素复制给堆顶
  // 同时减小堆的长度,模拟 出堆
  if (len > 1) {
    s[0] = s[len - 1]
    // 最后一个元素给堆顶之后,通过 下沉sink 保持最小堆的特性
    sink(0, len - 1)
  }
  len--
  return top
}
  • 💁 sink
function sink(idx, lastIdx) {
  while(idx <= lastIdx) {
    // 下沉索引位置与其左右子节点的值相比较,与较小值交换位置
    // 循环执行上述逻辑
    const rightIdx = (idx + 1) * 2
    const leftIdx = rightIdx - 1
    if (leftIdx <= lastIdx) {
      // 找到较小值
      let smallValueIdx = leftIdx
      if (rightIdx <= lastIdx && compare(smallValueIdx, rightIdx) > 0) {
        smallValueIdx = rightIdx
      }
      // 如果比较小值更小,交换位置
      if (compare(idx, smallValueIdx) > 0) {
        swap(idx, smallValueIdx)
        idx = smallValueIdx
      } else {
        break
      }
    } else {
      break
    }
  }
}
  • 工具函数
// 交换索引位置值
function swap(a, b) {
  let temp = s[a]
  s[a] = s[b]
  s[b] = temp
}

// 比较函数
function compare(a, b) {
  return s[a] - s[b]
}