最小的k个数

输入整数数组 arr ，找出其中最小的 k 个数。例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。

示例 1：

输入： arr = [3,2,1], k = 2

输出： [1,2] 或者 [2,1]

示例 2：

输入： arr = [0,1,2,1], k = 1

输出： [0]

限制：

• 0 <= k <= arr.length <= 10000
• 0 <= arr[i] <= 10000

/**
 * @param {number[]} arr
 * @param {number} k
 * @return {number[]}
 */
var getLeastNumbers = function (arr, k) {};

题目分析

虽然这题在 leetcode 上标注的「简单」，但是本题还是很有研究意义的。本文介绍了 3 种解法，时间复杂度依次降低，都是基于经典的算法或者数据结构。

解法 1: 直接排序

先说最简单、最直观的做法：直接排序。将数组按照从小到大的顺序排序，并且返回前 k 个数字。代码实现如下：

/**
 * @param {number[]} arr
 * @param {number} k
 * @return {number[]}
 */
var getLeastNumbers = function (arr, k) {
  return arr.sort((a, b) => a - b).slice(0, k);
};

使用高级排序（代码用的是快排），时间复杂度是O(NlogN)，空间复杂度是O(logN)。

解法 2: 最大堆

堆是一种非常常用的数据结构。最大堆的性质是：节点值大于子节点的值，堆顶元素是最大元素。利用这个性质，整体的算法流程如下：

• 创建大小为 k 的最大堆
• 将数组的前 k 个元素放入堆中
• 从下标 k 继续开始依次遍历数组的剩余元素：
  • 如果元素小于堆顶元素，那么取出堆顶元素，将当前元素入堆
  • 如果元素大于/等于堆顶元素，不做操作

由于堆的大小是 K，空间复杂度是O(K)，时间复杂度是O(NlogK)。

由于 JavaScript 中没有堆，所以需要手动实现。代码如下：

function swap(arr, i, j) {
  [arr[i], arr[j]] = [arr[j], arr[i]];
}

class MaxHeap {
  constructor(arr = []) {
    this.container = [];
    if (Array.isArray(arr)) {
      arr.forEach(this.insert.bind(this));
    }
  }

  insert(data) {
    const { container } = this;

    container.push(data);
    let index = container.length - 1;
    while (index) {
      let parent = Math.floor((index - 1) / 2);
      if (container[index] <= container[parent]) {
        break;
      }
      swap(container, index, parent);
      index = parent;
    }
  }

  extract() {
    const { container } = this;
    if (!container.length) {
      return null;
    }

    swap(container, 0, container.length - 1);
    const res = container.pop();
    const length = container.length;
    let index = 0,
      exchange = index * 2 + 1;

    while (exchange < length) {
      // 如果有右节点，并且右节点的值大于左节点的值
      let right = index * 2 + 2;
      if (right < length && container[right] > container[exchange]) {
        exchange = right;
      }
      if (container[exchange] <= container[index]) {
        break;
      }
      swap(container, exchange, index);
      index = exchange;
      exchange = index * 2 + 1;
    }

    return res;
  }

  top() {
    if (this.container.length) return this.container[0];
    return null;
  }
}

/**
 * @param {number[]} arr
 * @param {number} k
 * @return {number[]}
 */
var getLeastNumbers = function (arr, k) {
  const length = arr.length;
  if (k >= length) {
    return arr;
  }

  const heap = new MaxHeap(arr.slice(0, k));
  for (let i = k; i < length; ++i) {
    if (heap.top() > arr[i]) {
      heap.extract();
      heap.insert(arr[i]);
    }
  }
  return heap.container;
};

解法 3: 基于快速排序的 partition

解法 1 中使用了快速排序，但其实并需要对全部元素进行排序，题目只需要前 k 个元素。

回顾快速排序中的 partition 操作，可以将元素arr[0]放入排序后的正确位置，并且返回这个位置index。利用 partition 的特点，算法流程如下：

• 如果index = k，说明第 k 个元素已经放入正确位置，返回前 k 个元素
• 如果k < index，前 k 个元素在[left, index - 1]之间，缩小查找范围，继续查找
• 如果index < k，前 k 个元素在[index + 1, right] 之间，缩小查找范围，继续查找

为了方便理解，可以使用2, 8, 1, 1, 0, 11, -1, 0这个例子在纸上画一下过程。

代码实现如下：

/**
 *
 * @param {number[]} arr
 * @param {number} start
 * @param {number} end
 * @return {number}
 */
function partition(arr, start, end) {
  const k = arr[start];
  let left = start + 1,
    right = end;
  while (1) {
    while (left <= end && arr[left] <= k) ++left;
    while (right >= start + 1 && arr[right] >= k) --right;

    if (left >= right) {
      break;
    }

    [arr[left], arr[right]] = [arr[right], arr[left]];
    ++left;
    --right;
  }
  [arr[right], arr[start]] = [arr[start], arr[right]];
  return right;
}

/**
 * @param {number[]} arr
 * @param {number} k
 * @return {number[]}
 */
var getLeastNumbers = function (arr, k) {
  const length = arr.length;
  if (k >= length) return arr;
  let left = 0,
    right = length - 1;
  let index = partition(arr, left, right);
  while (index !== k) {
    if (index < k) {
      left = index + 1;
      index = partition(arr, left, right);
    } else if (index > k) {
      right = index - 1;
      index = partition(arr, left, right);
    }
  }

  return arr.slice(0, k);
};

时间复杂度是O(N)，空间复杂度是O(N)。