面试手撕🍣：用JS实现一个大根堆

raxskle
206 阅读5分钟

为了做力扣 215. 数组中的第 K 个最大元素 这道题目，而用 JS 手写了大根堆。所以看完本文不但能学会堆这个数据结构，还顺便刷了道力扣题。

先来看题

给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。

请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。

你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

示例 1:

输入: [3,2,1,5,6,4], k = 2
输出: 5

这道题明明可以用内置排序，再返回下标 k-1 元素

addd7e451f9d2c55da590e90197c8eb.png

简单一行代码就能过了！

27f2af474007a7d4e363e806c39c1e9.png

但要是面试碰上个面试官让手撕这题，总不能这么不给面试官面子吧

解法

为了实现时间复杂度为 O(n)得到第 k 大的元素，很自然可以想到堆排序这个方法。

大根堆或者优先队列这个数据结构，可以 O(1) 获取最大的元素。那么只需要获取 k 次，即可得到第 k 大的元素。而前提是把乱序的数组构建为一个大根堆，构建过程遍历数组一次就行，也就是时间复杂度 O(n)，看上去很完美，剩下的就是实现一个大根堆了。

堆

堆是一种特殊的树形数据结构，以大根堆为例，所有节点满足：父节点的值一定大于子节点的值，兄弟节点间的值没有大小约束。所以堆顶就是最大的元素，而小根堆则相反。

当新的元素新加入堆时，或者堆顶元素弹出时，都需要做平衡，使得仍然满足堆的要求。

image.png

实现

用数组模拟二叉树，将会让我们的实现代码更加简单，需要实现元素的插入和获取。

前置知识：用数组模拟完全二叉树时，如果根节点为 1，层序遍历给每个节点标号，对于某个节点 n，它的左子节点为 n*2，它的右子节点为 n*2+1

class MaxHeap {
  constructor() {
    // 用undefined占位，为了方便通过下标获取父元素，需要从1开始存根节点
    this.heap = [undefined];
  }

  // 返回节点的父节点下标
  getParentIndex(index) {
    return Math.floor(index / 2);
  }

  // 弹出堆顶，并平衡堆
  getMax() {}

  // 插入元素同时平衡堆
  insert(node) {}
}

insert()

先来实现 insert 方法，可以将数组通过 insert 元素不断地构建出一个完整的大根堆。

新加入的元素先 push 进 heap 数组的末尾（也就是完全二叉树的下一个坑位），然后不断地向上比较它与父节点的值，如果它比父节点更大，就把它和父节点交换。注意这里不需要管另一个兄弟节点，因为它的父节点原本就大于它的兄弟节点。

而交换也只需要将数组两个位置的元素交换即可。

  // 插入元素同时平衡堆
  insert(node) {
    this.heap.push(node);

    if (this.heap.length > 2) {
      // 需要平衡
      let currentIndex = this.heap.length - 1;
      while (
        currentIndex > 1 &&
        this.heap[this.getParentIndex(currentIndex)] < this.heap[currentIndex]
      ) {
        // 当父节点小于当前节点，需要将它们交换
        const parentIndex = this.getParentIndex(currentIndex);
        // 交换节点
        let temp = this.heap[parentIndex];
        this.heap[parentIndex] = this.heap[currentIndex];
        this.heap[currentIndex] = temp;
        currentIndex = parentIndex;
      }
    }
  }

getMax()

接下来实现弹出堆顶的方法。

弹出堆顶元素之后，向下寻找左右子节点中更大的一个 作为新的父节点，一直遍历整棵树，并将最后一个节点删除。

  // 返回堆顶
  getTop() {
    const top = this.heap[1];  // heap数组从第1位开始存
    if (this.heap.length < 2) {
      return top;
    } else if (this.heap.length === 2) {
      this.heap.pop();
      return top;
    }

    // 有子节点的情况
    let p = 1;
    while (
      this.heap[p * 2] !== undefined ||
      this.heap[p * 2 + 1] !== undefined
    ) {
      // 有左右任一子节点，需要找到更大的一个，作为新的父节点
      let nextIndex;
      if (this.heap[p * 2] === undefined) {
        // 左子节点为空
        nextIndex = p * 2 + 1;
      } else if (this.heap[p * 2 + 1] === undefined) {
        // 右子节点为空
        nextIndex = p * 2;
      } else if (this.heap[p * 2] < this.heap[p * 2 + 1]) {
        // 选择更大的节点
        nextIndex = p * 2 + 1;
      } else {
        nextIndex = p * 2;
      }
      this.heap[p] = this.heap[nextIndex]; // 取代父节点
      p = nextIndex; // 遍历下去
    }
    // 最后没有子节点的删除掉
    delete this.heap[p];

    return top;
  }

完整实现

class MaxHeap {
  constructor() {
    this.heap = [undefined]; // 占位
  }

  // 返回节点的父元素
  getParentIndex(index) {
    return Math.floor(index / 2);
  }

  // 返回堆顶
  getTop() {
    const top = this.heap[1];
    if (this.heap.length < 2) {
      return top;
    } else if (this.heap.length === 2) {
      this.heap.pop();
      return top;
    }

    // 有子节点的情况
    let p = 1;
    while (
      this.heap[p * 2] !== undefined ||
      this.heap[p * 2 + 1] !== undefined
    ) {
      // 有左右任一节点，需要平衡
      // 找到左右节点中更大的一个，作为新的节点
      let nextIndex;
      if (this.heap[p * 2] === undefined) {
        nextIndex = p * 2 + 1;
      } else if (this.heap[p * 2 + 1] === undefined) {
        nextIndex = p * 2;
      } else if (this.heap[p * 2] < this.heap[p * 2 + 1]) {
        // 选择更大的节点
        nextIndex = p * 2 + 1;
      } else {
        nextIndex = p * 2;
      }
      // 取代父节点
      this.heap[p] = this.heap[nextIndex];
      // 递归下去
      p = nextIndex;
    }
    // 最后没有子节点的删除掉
    delete this.heap[p];

    return top;
  }

  // 插入元素同时平衡堆
  insert(node) {
    this.heap.push(node);

    if (this.heap.length > 2) {
      // 需要平衡
      let currentIndex = this.heap.length - 1;
      while (
        currentIndex > 1 &&
        this.heap[this.getParentIndex(currentIndex)] < this.heap[currentIndex]
      ) {
        // 当父节点小于当前节点，需要将它们交换
        // 父亲节点index
        const parentIndex = this.getParentIndex(currentIndex);
        // 交换节点
        let temp = this.heap[parentIndex];
        this.heap[parentIndex] = this.heap[currentIndex];
        this.heap[currentIndex] = temp;
        currentIndex = parentIndex;
      }
    }
  }
}

const m = new MaxHeap();
const list = [3, 2, 1, 5, 6, 4];
for (let item of list) {
  m.insert(item);
}

console.log(m.getTop());
console.log(m.getTop());
console.log(m.getTop());
console.log(m.getTop());
console.log(m.getTop());
console.log(m.getTop());

输出

6
5
4
3
2
1

力扣 215. 数组中的第 K 个最大元素

最后，这道题的解法：

function findKthLargest(nums: number[], k: number): number {
  const maxHeap = new MaxHeap();
  // 构建堆
  for (const item of nums) {
    maxHeap.insert(item);
  }

  // 获取第k大的元素
  for (let i = 1; i < k; i++) {
    const res = maxHeap.getTop();
  }
  return maxHeap.getTop();
}

😉

avatar
文章
阅读
粉丝