【面试题】最小的k个数（三种解法）

题目描述

解法一：排序法

Java

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if(k ==0 || arr.length == 0) return new int[0];
        if(k == arr.length) return arr;
        Arrays.sort(arr); // 从小到大排序
        return Arrays.copyOfRange(arr, 0, k); // 返回前 k个数
    }
}

时间复杂度：O(n*logn)

解法二：快速选择法

Java

// 快速选择法
class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if(k ==0 || arr.length == 0) return new int[0];
        if(k == arr.length) return arr;
        int pivot_index = quick_select(arr, 0, arr.length-1, k); // 找到第 k小的元素所在位置，其前面的数字即为最小的前 k个数
        return Arrays.copyOfRange(arr, 0, pivot_index); // 返回前 k个数
    }

    public int partition(int[] nums, int left, int right, int pivot_index) {
        // 先取出枢纽元素，并将其交换到数组末尾
        int pivot = nums[pivot_index]; 
        nums[pivot_index] = nums[right];
        nums[right] = pivot;
        // 每次从未处理区间中取一个元素与pivot对比，如果小于pivot则将其放到pivot_index左边。
        int sorted_index = left;
        for(int i=left; i<right; i++) {
            if(nums[i] < pivot){
                int tmp = nums[i];
                nums[i] = nums[sorted_index];
                nums[sorted_index] = tmp;
                sorted_index++;
            }
        }
        // 循环结束后，将枢纽元素交换到合适的位置上
        int tmp = nums[sorted_index];
        nums[sorted_index] = nums[right];
        nums[right] = tmp;

        return sorted_index;
    }

    public int quick_select(int[] nums, int left, int right, int target) {
        if(left == right) return left; // 递归终止条件
        int pivot_index = left + (int)(Math.random()*(right-left+1)); // 每次随机选择一个位置作为分区的枢纽
        pivot_index = partition(nums, left, right, pivot_index); // 调用分区算法，确定排序后pivot_index合适的位置

        if(pivot_index == target) // 当前枢纽所在位置正好为目标位置
            return pivot_index;
        else if(target > pivot_index) // 如果目标位置大于当前枢纽所在位置，去右半边递归查找
            return quick_select(nums, pivot_index+1, right, target);
        else // 如果目标位置小于当前枢纽所在位置，去左半边递归查找
            return quick_select(nums, left, pivot_index-1, target);
    }
}

时间复杂度：O(n)
但是函数 Partition会改变数组中数字的顺序。

解法三：借助大顶堆（优先队列）

Java

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if(k == 0 || arr.length == 0) return new int[0];
        if(k == arr.length) return arr;
        // PriorityQueue 默认是自然顺序排序，也即小顶堆，如果实现大顶堆需要重写一下比较器
        Queue<Integer> heap = new PriorityQueue<>(new Comparator<Integer>(){
            @Override
			public int compare(Integer o1, Integer o2) {
				/*
				 *(1) 写成return o1.compareTo(o2) 或者 return o1-o2表示升序
				 *(2) 写成return o2.compareTo(o1) 或者return o2-o1表示降序
				 */
				return o2.compareTo(o1);
			}
        });

        for(int n:arr) {
            if(heap.size() < k) // 若目前堆的大小<K，直接将当前数字加入堆中。
                heap.offer(n);
            else if(heap.peek() > n) { // 当前数字小于堆顶元素才会入堆
                heap.poll(); // 删除堆顶即最大元素
                heap.offer(n);
            }           
        }

        // 将堆中的元素存入一个int数组并返回
        int[] res = new int[heap.size()];
        int i = 0;
        for(int num:heap) {
            res[i++] = num;
        }
        return res;
    }
}

时间复杂度：O(n*logk)，由于大顶堆需要实时维护前 k 小值，所以插入删除都是 O(logk) 的时间复杂度，最坏情况下数组里 n 个数都会插入，所以一共需要 O(nlogk) 的时间复杂度。
空间复杂度：O(k)，大顶堆最多 k个元素。

参考

https://leetcode-cn.com/problems/zui-xiao-de-kge-shu-lcof/solution/3chong-jie-fa-miao-sha-topkkuai-pai-dui-er-cha-sou/