leetcode-最小的k个数

 题目来自LeetCode，链接：面试题40. 最小的k个数。具体描述为：输入整数数组 arr ，找出其中最小的 k 个数。例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。限制：

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

 示例1：

输入：arr = [3,2,1], k = 2
输出：[1,2] 或者 [2,1]

 示例2：

输入：arr = [0,1,2,1], k = 1
输出：[0]

 一眼望去，就是一个可以用大顶堆来完成的一个题目嘛，先初始化一个大小为k的大顶堆来保存最小的k个数，对于前k个数先直接加入堆去，然后对于arr后面的数，判断是否小于堆顶元素，是的话堆顶元素出堆再将当前元素入堆，否则直接跳过。因为入堆操作的时间复杂度为 $O(logk)$，所以总时间复杂度为 $O(nlogk)$，空间复杂度为堆所占用空间即 $O(k)$。

 JAVA版代码如下：

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0) {
            return new int[0];
        }
        Queue mink = new PriorityQueue<Integer>(k, new Comparator<Integer>(){
            @Override
            public int compare(Integer i1, Integer i2) {
                return i2 - i1;
            }
        });
        for (int i = 0; i < k; ++i) {
            mink.offer(arr[i]);
        }
        for (int i = k; i < arr.length; ++i) {
            if (arr[i] < (int)mink.peek()) {
                mink.poll();
                mink.offer(arr[i]);
            }
        }
        int[] result = new int[k];
        Iterator iter = mink.iterator();
        int i = 0;
        while (iter.hasNext()) {
            result[i++] = (int)iter.next();
        }
        return result;
    }
}

 提交结果如下:

 看着时间不咋地，所以又想了一种空间换时间的方法，有点类似计数排序的思想，因为题目限定了数组arr中数的范围为[0, 10000]，所以可以用一个长为10001的数组record记录各个数出现次数，比如record[0]=1则说明0这个数出现了一次，那么最小的k个数就很容易得到了，即为record中前m个数（须满足前m-1个数的出现次数小于k，前m个数的出现次数大于等于k）。时间复杂度为 $O(n+k)$，空间复杂度为 $O(1)$。

 JAVA版代码如下：

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        int N = 10001;
        int[] record = new int[N];
        for (int i = 0; i < arr.length; ++i) {
            ++record[arr[i]];
        }
        int[] result = new int[k];
        for (int i = 0, idx = 0; i < N; ++i) {
            if (record[i] > 0) {
                for (int j = 0; j < record[i] && idx < k; ++j) {
                    result[idx++] = i;
                }
                if (idx == k) {
                    break;
                }
            }
        }
        return result;
    }
}

 提交结果如下:

 Python版代码如下：

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        N = 10001
        record = [0 for _ in range(N)]
        for a in arr:
            record[a] += 1
        idx = 0
        result = []
        for i in range(N):
            if record[i] > 0:
                for j in range(record[i]):
                    if idx < k:
                        result.append(i)
                        idx += 1
                    else:
                        break
                if idx == k:
                    break
        return result

 提交结果如下:

 然后又从评论区看到可以应用快排的想法，因为快排的划分函数partition每次执行完后都能将数组分成两个部分，小于分界值pivot的元素的都会被放到数组的左边，大于等于的都会被放到数组的右边，然后返回分界值的下标。也就是说在分界值左边包括分界值本身的M个数就一定是最小的M个数。那么定义函数 quick_select(arr, left, right, k) 表示划分数组 arr 的 [l,r] 部分，使前 k 小的数在数组的左侧，在函数里我们调用快排的划分函数，假设划分函数返回的下标是 middle（表示分界值 pivot 最终在数组中的位置），即 pivot 是数组中第 num = middle - left + 1 个数，那么一共会有三种情况：

• 如果 num == k，表示 pivot 及其左边的数就是最小的k个数，直接返回即可；

• 如果 num < k，表示第 k 小的数在 pivot 的右侧，因此递归调用 quick_select(arr, middle + 1, right, k - num)；

• 如果 num > k，表示第 k 小的数在 pivot 的左侧，递归调用 quick_select(arr, left, middle - 1, k)。

  因为第一次切分的时候需要遍历整个数组(0 ~ n)得到分界标j，下次切分会遍历数组(0~j-1)的元素或者(j+1～n)的元素，可以看作每次调用partition遍历的元素数目都是上一次遍历的1/2，总共需要 $n + n/2 + n/4 + \cdots + 1\approx 2n$, 因此时间复杂度是 $O(n)$。

 JAVA版代码如下：

class Solution {
    private static void swap(int[] arr, int i, int j) {
        int tmp = arr[i];
        arr[i] = arr[j];
        arr[j] = tmp;
    }

    private static int random_partition(int[] arr, int left, int right) {
        int idx = left + (int)Math.random() % (right - left + 1);
        swap(arr, idx, right);
        return partition(arr, left, right);
    }

    private static int partition(int[] arr, int left, int right) {
        int pivot = arr[right];
        int i = left - 1;
        for (int j = left; j < right; ++j) {
            if (arr[j] < pivot) {
                swap(arr, ++i, j);
            }
        }
        swap(arr, ++i, right);
        return i;
    }

    /*
    private void quickSort(int[] arr, int left, int right) {
        if (left >= right) {
            return;
        }
        int middle = random_partition(arr, left, right);
        quickSort(arr, left, middle - 1);
        quickSort(arr, middle + 1, right);
    }
    */

    private void quickSelect(int[] arr, int left, int right, int k) {
        if (left >= right) {
            return;
        }
        int middle = random_partition(arr, left, right);
        int num = middle - left + 1;
        if (k == num) {
            return;
        }
        if (k < num) {
            quickSelect(arr, left, middle - 1, k);
        }
        else {
            quickSelect(arr, middle + 1, right, k - num);
        }
    }

    public int[] getLeastNumbers(int[] arr, int k) {        
        int[] result = new int[k];
        //quickSort(arr, 0, arr.length - 1);
        quickSelect(arr, 0, arr.length - 1, k);
        for (int i = 0; i < k; ++i) {
            result[i] = arr[i];
        }
        return result;
        
    }
}

 提交结果如下: