本问题对应的 leetcode 原文链接：剑指 Offer 40. 最小的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

解题思路

视频讲解直达： 本题视频讲解

代码实现

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if(arr == null || arr.length == 0 || k == 0){
            return new int[0];
        }

        return quickFind(arr, 0, arr.length - 1, k);
    }

    int[] quickFind(int[] arr, int left, int right, int k){
        int i = partition(arr, left, right);
        // 之所以需要 i+1，是因为下标从 0 开始，0～i之间一共有 i+1个数
        if(i + 1 == k){
            return Arrays.copyOf(arr, k);
        }

        if(i + 1 > k){
            return quickFind(arr, 0, i - 1, k);
        } else {
            return quickFind(arr, i + 1, right, k);
        }
    }
    // 找出pivot的下标以及使小于等于pivot在左边，大于等于的在右边
    int partition(int[] arr, int left, int right){
        int pivot = arr[left];

        int i = left + 1;
        int j = right;

        while(i < j){
            while(i <= j && arr[i] <= pivot) i++;
            while(i <= j && arr[j] >= pivot) j--;
            if(i >= j){
                break;
            }

            int temp = arr[i];
            arr[i] = arr[j];
            arr[j] = temp;
        }

        arr[left] = arr[j];
        arr[j] = pivot;
        return j;
    }
}

时间复杂度：O(n)
空间复杂度：O(logn)，平均递归调用深度

Python

class Solution(object):
    def getLeastNumbers(self, arr, k):
        """
        :type arr: List[int]
        :type k: int
        :rtype: List[int]
        """
        if not arr or len(arr) == 0 or k == 0:
            return []

        return self.quickFind(arr, 0, len(arr) - 1, k)

    def quickFind(self, arr, left, right, k):
        i = self.partition(arr, left, right)
        if i + 1 == k:
            return arr[:k]
        elif i + 1 > k:
            return self.quickFind(arr, left, i - 1, k)
        else:
            return self.quickFind(arr, i + 1, right, k)

    def partition(self, arr, left, right):
        pivot = arr[left]

        i = left + 1
        j = right

        while i <= j:
            while i <= j and arr[i] <= pivot:
                i += 1
            while i <= j and arr[j] >= pivot:
                j -= 1
            if i >= j:
                break

            temp = arr[i]
            arr[i] = arr[j]
            arr[j] = temp

        arr[left] = arr[j]
        arr[j] = pivot
        return j

C++

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        if(arr.empty() || k == 0){
            return {};
        }

        return quickFind(arr, 0, arr.size() - 1, k);
    }

    vector<int> quickFind(vector<int>& arr, int left, int right, int k){
        int i = partition(arr, left, right);
        if(i + 1 == k){
            return vector<int>(arr.begin(), arr.begin() + k);
        } else if(i + 1 > k){
            return quickFind(arr, left, i - 1, k);
        } else {
            return quickFind(arr, i + 1, right, k);
        }
    }

    int partition(vector<int>& arr, int left, int right){
        int pivot = arr[left];

        int i = left + 1;
        int j = right;

        while(i <= j){
            while(i <= j && arr[i] <= pivot) i++;
            while(i <= j && arr[j] >= pivot) j--;
            if(i >= j){
                break;
            }

            int temp = arr[i];
            arr[i] = arr[j];
            arr[j] = temp;
        }

        arr[left] = arr[j];
        arr[j] = pivot;
        return j;
    }
};

Go

func getLeastNumbers(arr []int, k int) []int {
    if len(arr) == 0 || k == 0{
        return []int{}
    }

    return quickFind(arr, 0, len(arr) - 1, k)
}

func quickFind(arr []int, left int, right int, k int) []int {
    i := partition(arr, left, right)
    if i + 1 == k{
        return arr[:k]
    } else if i + 1 > k{
        return quickFind(arr, left, i - 1, k)
    } else {
        return quickFind(arr, i + 1, right, k)
    }
}

func partition(arr []int, left int, right int) int {
    pivot := arr[left]

    i := left + 1
    j := right

    for i <= j {
        for i <= j && arr[i] <= pivot{
            i++
        }
        for i <= j && arr[j] >= pivot {
            j--
        }
        if i >= j {
            break
        }

        arr[i], arr[j] = arr[j], arr[i]
    }

    arr[left], arr[j] = arr[j], pivot
    return j
}

JS

/**
 * @param {number[]} arr
 * @param {number} k
 * @return {number[]}
 */
var getLeastNumbers = function(arr, k) {
  if (arr === null || arr.length === 0 || k === 0) {
    return new Array(0);
  }

  return quickFind(arr, 0, arr.length - 1, k);
};

function quickFind(arr, left, right, k) {
  var i = partition(arr, left, right);
  // 之所以需要 i+1，是因为下标从 0 开始，0～i之间一共有 i+1个数
  if (i + 1 === k) {
    return arr.slice(0, k);
  }

  if (i + 1 > k) {
    return quickFind(arr, 0, i - 1, k);
  } else {
    return quickFind(arr, i + 1, right, k);
  }
}

// 找出pivot的下标以及使小于等于pivot在左边，大于等于的在右边
function partition(arr, left, right) {
  var pivot = arr[left];

  var i = left + 1;
  var j = right;

  while (i < j) {
    while (i <= j && arr[i] <= pivot) i++;
    while (i <= j && arr[j] >= pivot) j--;
    if (i >= j) {
      break;
    }

    var temp = arr[i];
    arr[i] = arr[j];
    arr[j] = temp;
  }

  arr[left] = arr[j];
  arr[j] = pivot;
  return j;
}

shuaidi永久会员

收藏 海报挣佣金 推广链接