React 优先级队列小顶堆的简单实现

这篇文章是博主最近在阅读 React 源码过程中关于优先级队列数据结构实现的一些提炼

我们都知道 React Schedule 中的优先级队列，本身采用小顶堆来保证队列的第一个任务是优先级最高的，并且获取优先级最高的任务的时间复杂度为 O(1)。所以今天写了个简化版的小顶堆的实现，方便大家理解。

首先优先级队列有两种操作，一种是添加任务，另外一种是移除优先级最高的任务（第一个任务）

const heap = [];
// 添加任务
const push = (val) => {
  // 先将任务添加到队列中
  heap.push(val); 
  // 对当前任务节点进行向上调整，具体实现看后面
  siftUp(heap, val, heap.length);
};

// 移除优先级最高的任务
const pop = () => {
  // 拿到最后一个任务放到第一个任务来的位置上来
  const last = heap.pop();
  heap[0] = last;
  // 对当前任务节点进行向下调整，具体实现看后面
  siftDown(heap, last, 0);
};

// 获取优先级最高的任务
const peek = () => {
  return heap.length ? heap[0] : null;
};

当我们往队列中去添加新的任务的时候，那么这个任务是肯定会在队尾。但由于小顶堆的特性，我们需要保证第一个任务优先级要最高，所以需要拿新任务和他的所有父节点去进行对比。

const siftUp = (heap, node, i) => {
  // 需要把节点下标记录下来，
  let index = i - 1;
  while (index > 0) {
    // 拿到 push 进来的节点的父节点
    const parentIndex = index >>> 1; // 等价于 Math.floor((index - 1) / 2)
    const parent = heap[parentIndex];
    if (parent > node) {
      // 如果父节点大于当前节点，说明当前节点需要和父节点交换位置
      heap[parentIndex] = node;
      heap[index] = parent;
      index = parentIndex;
    } else {
      return;
    }
  }
};

移除操作也是同理，移除一个高优先级任务之后，会把队尾的任务提高最顶部来，为了保证小顶堆的特性，所以需要拿这个任务和其子节点进行对比。

const siftDown = (heap, node, i) => {
  let index = i; // 当前需要处理节点的下标
  const idx = heap.length >>> 1; // 拿到最深的非叶子结点的下标，由于小顶堆的特性，不需要到叶子节点去
  while (index < idx) {
    let leftIndex = 2 * index + 1;
    let left = heap[leftIndex];
    let rightIndex = leftIndex + 1;
    let right = heap[rightIndex];

    // 如果当前节点，比左子节点要大，说明需要调整
    if (node > left) {
      // 如果左子节点比右子节点要大，说明右子节点是目前最小的，应该和当前节点交换位置
      if (left > right) {
        heap[rightIndex] = node;
        heap[index] = right;
        index = rightIndex;
      } else {
        // 如果左子节点比右子节点要小，说明左子节点是目前最小的，应该和当前节点交换位置
        heap[leftIndex] = node;
        heap[index] = left;
        index = leftIndex;
      }
    } else if (node > right) {
      // 如果当前节点，比右子节点要大，说明需要交换位置
      heap[rightIndex] = node;
      heap[index] = right;
      index = rightIndex;
    } else {
      return;
    }
  }
};

写完了上面的代码之后，我们可以测一下

push(5);
push(7);
push(10);
push(8);
push(2);
push(6);
console.log(heap); // [ 2, 5, 6, 8, 10, 7 ]

pop();
console.log(heap); // [ 5, 7, 6, 8, 10 ]

完整的代码如下：

const heap = [];

// 添加任务
const push = (val) => {
  // 先将任务添加到队列中
  heap.push(val); 
  // 对当前任务节点进行向上调整，具体实现看后面
  siftUp(heap, val, heap.length);
};

// 移除优先级最高的任务
const pop = () => {
  // 拿到最后一个任务放到第一个任务来的位置上来
  const last = heap.pop();
  heap[0] = last;
  // 对当前任务节点进行向下调整，具体实现看后面
  siftDown(heap, last, 0);
};

// 获取优先级最高的任务
const peek = () => {
  return heap.length ? heap[0] : null;
};

const siftUp = (heap, node, i) => {
  // 需要把节点下标记录下来，
  let index = i - 1;
  while (index > 0) {
    // 拿到 push 进来的节点的父节点
    const parentIndex = index >>> 1; // 等价于 Math.floor((index - 1) / 2)
    const parent = heap[parentIndex];
    if (parent > node) {
      // 如果父节点大于当前节点，说明当前节点需要和父节点交换位置
      heap[parentIndex] = node;
      heap[index] = parent;
      index = parentIndex;
    } else {
      return;
    }
  }
};

const siftDown = (heap, node, i) => {
  let index = i; // 当前需要处理节点的下标
  const idx = heap.length >>> 1; // 拿到最深的非叶子结点的下标
  while (index < idx) {
    let leftIndex = 2 * index + 1;
    let left = heap[leftIndex];
    let rightIndex = leftIndex + 1;
    let right = heap[rightIndex];

    // 如果当前节点，比左子节点要大，说明需要调整
    if (node > left) {
      // 如果左子节点比右子节点要大，说明右子节点是目前最小的，应该和当前节点交换位置
      if (left > right) {
        heap[rightIndex] = node;
        heap[index] = right;
        index = rightIndex;
      } else {
        // 如果左子节点比右子节点要小，说明左子节点是目前最小的，应该和当前节点交换位置
        heap[leftIndex] = node;
        heap[index] = left;
        index = leftIndex;
      }
    } else if (node > right) {
      // 如果当前节点，比右子节点要大，说明需要交换位置
      heap[rightIndex] = node;
      heap[index] = right;
      index = rightIndex;
    } else {
      return;
    }
  }
};

push(5);
push(7);
push(10);
push(8);
push(2);
push(6);
console.log(heap);

pop();
console.log(heap);

参考：React 源码 - 最小堆实现