最小的k个数(TopN 问题)

Offer 驾到，掘友接招！我正在参与2022春招打卡活动，点击查看活动详情。

一、题目描述：

输入整数数组，找出其中最小的 k 个数。这是典型的 TopN 问题。

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

二、思路分析：

根据题目的要求，找到数组中的最小的k个数，那么我们可以先把数组从小到大进行排序，然后取出前k个数即可，于是就写出了下面：

function getLeastNumbers (arr, k) {
  return arr.sort((a, b) => a - b).slice(0, k)
}

时间复杂度：O(nlogn)。虽然使用数组内置的方法可以实现，但是复杂度较高，再想想其他办法。上面是排序所有数，其实我们只需要排序k个数就可以了。

我们用一个栈实时维护数组的前 k 小值。首先将前 k 个数插入栈中，随后从第 k+1 个数开始遍历，如果当前遍历到的数比栈顶的数要小，就把栈顶的数弹出，再插入当前遍历到的数。

function getLeastNumbers (arr, k) {
  let queue = []
  for (let i = 0; i < k; i++) {
    queue.push(arr[i])
  }
  queue = queue.sort((a, b) => a - b)
  for (let i = k; i < arr.length; i++) {
    if (queue[queue.length - 1] > arr[i]) {
      queue.pop()
      queue.push(arr[i])
      queue = queue.sort((a, b) => a - b)
    }
  }
  return queue
}

console.log(getLeastNumbers([6,2,3,4,5], 2))

时间复杂度：O(nlogk)，由于实时维护前 k 小值。

首先将前 k 个数插入数组中，[2,6]
从3开始遍历
当遍历到3，移除6加入3，数组为[2,3]
当遍历到4，不插入
当遍历到5，不插入
当前栈实时维护数组的前 k 小值

注意，每次加入元素之后需要进行重新排序，因为每次只检查栈顶元素，具体元素插入的位置需要重新排序。

对栈的两种主要操作是将一个元素压入栈和将一个元素弹出栈。入栈使用 push() 方法，出栈使用 pop() 方法。

方法二是类似于快排的思想，但是实现起来难度较大。时间复杂度 O(N)

function getLeastNumbers (arr, k) {
  if (k >= arr.length) {
    return arr
  }
  return quickSort(arr, k, 0, arr.length - 1)
}

function quickSort (arr, k, l, r) {
  let i = l
  let j = r
  while (i < j) {
    while (i < j && arr[j] >= arr[l]) {
      j--
    }
    while (i < j && arr[i] <= arr[l]) {
      i++
    }
    swap(arr, i, j)
  }
  swap(arr, i, l)
  if (i > k) {
    return quickSort(arr, k, l, i - 1)
  }
  if (i < k) {
    return quickSort(arr, k, i + 1, r)
  }
  return arr.slice(0, k)
}

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

console.log(getLeastNumbers([1,2,3,4,5], 2))

四、总结：

这道题是一个经典的 Top K 问题，是面试中的常客。Top K 问题有两种不同的解法，一种解法使用栈（维护已排序数组），另一种解法使用类似快速排序的分治法。这两种方法各有优劣，最好都掌握。