剑指offer - 最小的k个数 - JavaScript

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

限制:

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

# 题目分析

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

# 解法 1: 直接排序

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

// ac地址:https://leetcode-cn.com/problems/zui-xiao-de-kge-shu-lcof/
// 原文地址:https://xxoo521.com/2020-02-21-least-nums/

/**
 * @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 中没有堆,所以需要手动实现。代码如下:

// ac地址:https://leetcode-cn.com/problems/zui-xiao-de-kge-shu-lcof/
// 原文地址:https://xxoo521.com/2020-02-21-least-nums/
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这个例子在纸上画一下过程。

代码实现如下:

// ac地址:https://leetcode-cn.com/problems/zui-xiao-de-kge-shu-lcof/
// 原文地址:https://xxoo521.com/2020-02-21-least-nums/

/**
 *
 * @param {number[]} arr
 * @param {number} start
 * @param {number} end
 * @return {number}
 */
function partiton(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 = partiton(arr, left, right);
  while (index !== k) {
    if (index < k) {
      left = index + 1;
      index = partiton(arr, left, right);
    } else if (index > k) {
      right = index - 1;
      index = partiton(arr, left, right);
    }
  }

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

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

# 总结

不得不说 Leetcode 平台还是很准确的。下面第一张是解法 3 的运行结果,第二张是解法 2 的运行结果,解法 3 的时间效率更高。

解法 2 中实现的堆不是最优的方法,有 2 点可以优化:

  • swap 交换换成元素赋值。
  • 根绝数组创建堆的时候,可以从第一个非叶节点向上执行 shift_up 操作。这个过程是heapify

不过本题考察重点不在这,所以就随手写了下堆的代码(主要是看的明白)。有兴趣可以看一下 C++标准库的实现。

# 更多资料

整理不易,若对您有帮助,请给个「关注+点赞」,您的支持是我更新的动力 👇