[leetcode]面试题40. 最小的k个数

• 个人博客：https://javaniuniu.com/
• 难度：简单
• 本题涉及算法： 小根堆 大根堆 快排``二叉树
• 思路：小根堆 大根堆 快排``二叉树
• 类似题型：

---

题目 面试题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]

方法一 暴力

java

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        for (int i=0;i<arr.length-1;i++) {
                if (arr[i]>arr[i+1]){
                    int temp = arr[i];
                    arr[i] = arr[i+1];
                    arr[i+1] = temp;
                }
                for (int j=i;j>0;j--) {
                    if(arr[j]<arr[j-1]) {
                        int temp = arr[j];
                        arr[j] = arr[j-1];
                        arr[j-1] = temp;
                    }
                }
            }
            int[] res = new int[k];
            for (int i = 0;i<k;i++) {
                res[i] = arr[i];
            }
            return res;
    }
}

python

class Solution(object):
    def getLeastNumbers(self, arr, k):
        """
        :type arr: List[int]
        :type k: int
        :rtype: List[int]
        """
        for i in range(len(arr)-1):
            if arr[i]> arr[i+1]:
                sort(arr,i)
            for j in range(i,0,-1):
                if arr[j] < arr[j-1]:
                    lowsort(arr,j)
        return arr[:k]

def sort(arr,i):
    temp = arr[i]
    arr[i] = arr[i+1]
    arr[i+1] = temp
    return arr

def lowsort(arr,j):
    temp = arr[j]
    arr[j] = arr[j-1]
    arr[j-1] = temp
    return arr

方法二 小根堆

/**
   * 小根堆
   * @param arr
   * @param k
   * @return
   */
  public int[] getLeastNumbersByStack(int[] arr,int k) {
        if (k == 0 || arr.length == 0) {
            return new int[0];
        }
        // 默认是小根堆，实现大根堆需要重写一下比较器。
        Queue<Integer> pq = new PriorityQueue<>((v1, v2) -> v2 - v1); // 目的是最大数在第一位
        for(int num :arr) {
        if(pq.size()<k)
            pq.offer(num);
        else if(pq.peek()>num) {
            pq.poll();
            pq.offer(num);
        }
      }
      // 返回堆中的元素
      int[] res = new int[pq.size()];
      int idx = 0;
      for(int num: pq) {
          res[idx++] = num;
      }
      return res;
  }

方法三 大根堆

我们用一个大根堆实时维护数组的前 kk 小值。首先将前 kk 个数插入大根堆中，随后从第 k+1k+1 个数开始遍历，如果当前遍历到的数比大根堆的堆顶的数要小，就把堆顶的数弹出，再插入当前遍历到的数。最后将大根堆里的数存入数组返回即可

class Solution:
    def getLeastNumbers(self, nums: List[int], k: int) -> List[int]:
        if k == 0: return []

        n, opposite = len(nums), [-1 * x for x in nums[:k]]
        heapq.heapify(opposite)
        for i in range(k, len(nums)):
            if -opposite[0] > nums[i]:
                # 维持堆大小不变
                heapq.heappop(opposite)
                heapq.heappush(opposite, -nums[i])
        return [-x for x in opposite]

方法三 二叉树

/**
     * 二叉树
     * @param arr
     * @param k
     * @return
     */
public int[] getLeastNumbersByTree(int[] arr, int k) {
        if (k == 0 || arr.length == 0) {
            return new int[0];
        }
        // TreeMap的key是数字, value是该数字的个数。
        // cnt表示当前map总共存了多少个数字。
        TreeMap<Integer, Integer> map = new TreeMap<>();
        int cnt = 0;
        for (int num: arr) {
            // 1. 遍历数组，若当前map中的数字个数小于k，则map中当前数字对应个数+1
            if (cnt < k) {
                map.put(num, map.getOrDefault(num, 0) + 1);
                cnt++;
                continue;
            }
            // 2. 否则，取出map中最大的Key（即最大的数字), 判断当前数字与map中最大数字的大小关系：
            //    若当前数字比map中最大的数字还大，就直接忽略；
            //    若当前数字比map中最大的数字小，则将当前数字加入map中，并将map中的最大数字的个数-1。
            Map.Entry<Integer, Integer> entry = map.lastEntry();
            if (entry.getKey() > num) {
                map.put(num, map.getOrDefault(num, 0) + 1);
                if (entry.getValue() == 1) {
                    map.pollLastEntry();
                } else {
                    map.put(entry.getKey(), entry.getValue() - 1);
                }
            }

        }
}

方法五 数据范围有限时直接计数排序就行了：O(N)

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0 || arr.length == 0) {
            return new int[0];
        }
        // 统计每个数字出现的次数
        int[] counter = new int[10001];
        for (int num: arr) {
            counter[num]++;
        }
        // 根据counter数组从头找出k个数作为返回结果
        int[] res = new int[k];
        int idx = 0;
        for (int num = 0; num < counter.length; num++) {
            while (counter[num]-- > 0 && idx < k) {
                res[idx++] = num;
            }
            if (idx == k) {
                break;
            }
        }
        return res;
    }
}

方法六 排序

# 这个方法很快，总感觉有作弊嫌疑
class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        arr.sort()
        return arr[:k]