最小堆：TopK问题

在解决TopK问题时，综合各方面考虑，最小堆方案比较好，其时间复杂度为O(n*lgK)，空间复杂度为O(K)。

由于空间复杂度取决于K值，因此，可应用于大量流数据的TopK问题（无法将全部数据载入内存的情况）。

思路：

（1）构造一个大小为K的最小堆，用于存储当前TopK元素，堆顶为TopK元素中最小的元素，即第K个元素；

（2）每读入一个元素，与堆顶元素比较，若大于堆顶元素，则替换堆顶元素，并将新的堆顶元素进行排序。

代码如下：

public class TopK {

  private int k;
  private List<Integer> heap;  // 存放最小堆数据

  public TopK(int k) {
    this.k = k;
    this.heap = new ArrayList<>(k);
  }

  /**
   * 输入/读入一个元素
   */
  public void input(int num) {
    if (heap.size() < k) {  // 最小堆未满：尾部添加元素，并通过上浮法对其排序。
      heap.add(num);
      swim();
    } else if (num > heap.get(0)) {  // 最小堆已满 且 当前元素大于堆顶元素：替换堆顶元素，并通过下沉法对其排序。
      heap.set(0, num);
      sink();
    }
  }

  /**
   * 上浮法排序：将新的尾部元素进行排序
   */
  private void swim() {
    int index = heap.size() - 1;
    int parentIndex = (index - 1) / 2;
    while (index > 0 && heap.get(index) < heap.get(parentIndex)) {
      Collections.swap(heap, index, parentIndex);
      index = parentIndex;
      parentIndex = (index - 1) / 2;
    }
  }

  /**
   * 下沉法排序：将堆顶元素进行排序
   */
  private void sink() {
    int index = 0;
    while (index < heap.size()) {
      int minValue = heap.get(index), swapIndex = -1;
      int leftIndex = 2 * index + 1, rightIndex = 2 * index + 2;
      if (leftIndex < heap.size() && heap.get(leftIndex) < minValue) {
        minValue = heap.get(leftIndex);
        swapIndex = leftIndex;
      }
      if (rightIndex < heap.size() && heap.get(rightIndex) < minValue) {
        minValue = heap.get(rightIndex);
        swapIndex = rightIndex;
      }
      if (swapIndex > -1) {
        Collections.swap(heap, index, swapIndex);
        index = swapIndex;
      } else {
        break;
      }
    }
  }

  /**
   * 获取第K个值：即堆顶元素
   */
  public int getValueAtK() {
    return heap.get(0);
  }

  /**
   * 获取TopK值列表（无序）
   */
  public List<Integer> getTopK() {
    return new ArrayList<>(heap);  // 深拷贝：避免外部修改影响内部数据
  }

  /**
   * 获取TopK值列表（有序）
   */
  public List<Integer> getSortedTopK() {
    List<Integer> retList = new ArrayList<>(heap);  // 深拷贝：避免修改影响原堆数据
    Collections.sort(retList);
    Collections.reverse(retList);
    return retList;
  }
}

相关文章

《全排列（Java）》

《位运算：减法与补码》

《异或（^）的性质与应用》

《图解：常用排序算法（冒泡、选择、插入、希尔、快速、归并、堆）》

《回溯算法（试探算法）》

《动态规划：鸡蛋掉落》

《动态规划：单词拆分》

《状态机：只出现一次的数字II》

《链表：快慢指针》