base算法与数据结构【Q288】如何求数组中的 TOP k

如何求数组中的 TOP k

更多描述 求数组中的前 N 个最大的数

Issue 欢迎在 Gtihub Issue 中回答此问题: Issue 290

Author 回答者: shfshanyue

  1. 取数组中前 k 个数做小顶堆,堆化
  2. 数组中的其它数逐一与堆顶元素比较,若大于堆顶元素,则插入该数

时间复杂度 O(nlg(k))

Author 回答者: manyyuri

实现一个优先队列类,默认大顶堆,传入(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);