[路飞] 最小 k 个数

题目描述

设计一个算法，找出数组中最小的k个数。以任意顺序返回这k个数均可。

解题思路

算法

快排 partition

思路

我们通过快排 partition 的思想，把数组划分成最左面为 k 长度的部分，那么根据快排可得，划分的各个部分是相对有序的，因此这部分最左边的，长度为 k 的部分，是最小 k 个数，我们 slice 之后返回即可

过程

利用快排的 partition 思想，找到一个 pivot（这里是随机数），将数组从 l(default to 0) 至 r(default to length - 1) 进行划分，要注意一点的是，我们只会对左区间进行分割，因此要找的是最小 k。

每次 partition 完成之后，记录下来区间的最右边索引 pos，此时：

pos - l + 1 是我已经划分的区间长度，

根据不同的划分情况，我们需要在下次 partition 中传不同的 l, r

如果 > k，改变 r 为 pos - 1，因此划分出来的长度已经超过了 k，所以我们减小范围，在更小的区间内划分

如果 < k，改变 l 为 pos + 1, 长度不够，还要继续划分，所以剩余需要划分的长度为 k - (pos - l + 1）

代码

/**
 * @param {number[]} arr
 * @param {number} k
 * @return {number[]}
 */
var smallestK = function (arr, k) {
    select(arr, 0, arr.length - 1, k)
    return arr.slice(0, k)
};

function select(arr, l, r, k) {
    if (l >= r) return

    const pos = partition(arr, l, r)
    const num = pos - l + 1
    if (num === k) return
    else if (num > k) {
        return select(arr, l, pos - 1, k)
    } else {
        return select(arr, pos + 1, r, k - num)
    }
}

function partition(arr, l, r) {
    const index = parseInt(Math.random() * (r - l + 1)) + l
    ;[arr[r], arr[index]] = [arr[index], arr[r]]
    const pivot = arr[r]

    let i = l - 1
    for (let j = l; j < r; j++) {
        if (arr[j] <= pivot) {
            i++
            ;[arr[i], arr[j]] = [arr[j], arr[i]]
        }
    }

    ;[arr[r], arr[i + 1]] = [arr[i + 1], arr[r]]
    return i + 1
}