DeTechn Blog

如何求数组中的 TOP k

取数组中前 k 个数做小顶堆,堆化
数组中的其它数逐一与堆顶元素比较,若大于堆顶元素,则插入该数
时间复杂度 O(nlg(k))
实现一个优先队列类,默认大顶堆,传入(x,y)=>x>y 比较函数则为小顶堆。 首先将前 k 个数 insert,之后的数字如果大于栈顶元素(栈顶元素为堆中最小值),delTop 删除栈顶元素,然后 insert 该数。 维护一个 size 为 k 的小顶堆。 最后堆中元素即为数组中最大的 k 个元素。

class PriorityQueue {
  constructor(
    // 默认大顶堆
    cmp = (x, y) => {
      return x < y;
    }
  ) {
    this.queue = [];
    this.N = 0;
    this.cmp = (i, j) => {
      return cmp(this.queue[i], this.queue[j]);
    };
    this.parent = function (x) {
      return Math.floor(x / 2);
    };
    this.left = function (x) {
      return x * 2;
    };
    this.right = function (x) {
      return x * 2 + 1;
    };
    this.exch = (x, y) => {
      let temp = this.queue[x];
      this.queue[x] = this.queue[y];
      this.queue[y] = temp;
    };
  }
  swim(k) {
    const { cmp, parent, exch } = this;
    while (k > 1 && cmp(parent(k), k)) {
      exch(parent(k), k);
      k = parent(k);
    }
  }
  sink(k) {
    const { left, right, exch, N, cmp } = this;
    while (left(k) <= N) {
      let older = left(k);
      if (right(k) <= N && cmp(older, right(k))) {
        older = right(k);
      }
      if (cmp(older, k)) break;
      exch(older, k);
      k = older;
    }
  }
  top() {
    return this.queue[1];
  }
  delTop() {
    const { queue, exch } = this;
    let top = queue[1];
    exch(1, this.N);
    queue.pop();
    this.N--;
    this.sink(1);
    return top;
  }
  insert(x) {
    this.N++;
    this.queue[this.N] = x;
    this.swim(this.N);
  }
  size() {
    return this.N;
  }
}

function TopK(arr, k) {
  // 传入cmp,设置为小顶堆
  let pq = new PriorityQueue((x, y) => x > y);
  let i = 0;
  for (i = 0; i < k; i++) {
    pq.insert(arr[i]);
  }
  for (; i < arr.length; i++) {
    if (arr[i] > pq.top()) {
      pq.delTop();
      pq.insert(arr[i]);
    }
  }
  console.log("TOP K应为:", arr.sort((a, b) => b - a).splice(0, k));

  console.log("求出的TOP k:");
  // 排序,整理,方便对照
  pq.queue.sort((a, b) => b - a).pop();
  console.log(pq.queue);
}
let arr = [10, 15, 2, 6, 4, 5, 7, 3, 6, 14, 3, 12, 14, 13, 16, 1, 8];
TopK(arr, 10);

当前页面是本站的「Google AMP」版。查看和发表评论请点击:完整版 »