Java 实现最小的K个数(类似求topN)

输入n个整数，找出其中最小的K个数。例如输入4,5,1,6,2,7,3,8这8个数字，则最小的4个数字是1,2,3,4。

代码

解法一

先对数组排序，再取出k个元素，时间复杂度等于排序方法的复杂度，会改变原数组，但是取出的k个元素是有序的

    public static int[] findLeastK(int[] array, int k) {
        if (array == null || array.length == 0 || k <= 0) {
            return null;
        }
        sort(array, 0, array.length - 1);
        int[] result = new int[k];
        for (int i = 0; i < k; i++) {
            result[i] = array[i];
        }
        return result;
    }

    /**
     * 快排
     * @param array
     * @param lo
     * @param hi
     */
    public static void sort(int[] array,int lo ,int hi){
        if(lo>=hi){
            return ;
        }
        int index=partition(array,lo,hi);
        sort(array,lo,index-1);
        sort(array,index+1,hi);
    }

    public static int partition(int []array,int start,int end){
        // 固定的切分方式
        int key = array[start];
        while(start < end){
            // 从后半部分向前扫描
            while(end > start && array[end] >= key){
                end--;
            }
            array[start] = array[end];
            // 从前半部分向后扫描
            while(start < end && array[start] < key){
                start++;
            }
            array[end]=array[start];
        }
        array[start]=key;
        return end;
    }

    public static void main(String[] args) {
        int[] array = {4, 5, 1, 6, 2, 7, 3, 8};
        int[] result = findLeastK(array, 4);
        for (int i : result) {
            System.out.print(i + " ");
        }
    }

输出是 1 2 3 4

解法二

求数组的最小K个数，可以理解为求topK(最小堆)，采用最大堆，适用于大数据量小K的情况，O(n*logk)，不会改变原数组，但是结果是无序的

    /**
     * 最大堆
     */
    static class MaxHeap {
        // 堆的存储结构 - 数组，father=data[2 * i],left=data[2 * i +1], right=data[2 * i + 2]
        private int[] data;

        /**
         * 将一个数组传入构造方法，并转换成一个小根堆
         * @param data
         */
        public MaxHeap(int[] data) {
            this.data = data;
            buildMaxHeap();
        }

        /**
         * 将数组转换成最大堆
         * 数组中，父节点对应的位置是i*2，左子节点对应的是i * 2 + 1， 右子节点对应的是i * 2 + 2
         */
        private void buildMaxHeap() {
            // 完全二叉树只有数组下标小于或等于 (data.length) / 2 - 1 的元素有孩子结点，遍历这些结点。
            // 比如{0， 1， 2， 3， 4}，(data.length) / 2 - 1 = 1，其左右子节点是3和4，而2,3,4没有子节点
            for (int i = (data.length) / 2 - 1; i >= 0; i--) {
                // 对有子结点的元素heapify
                heapify(i);
            }
        }

        /**
         * 调整节点i及其子节点的位置以符合最大堆的要求，即root是最大值
         * 如果是求数组最大的K个数，则使用最小堆，即左右子节点比root大的调换位置即可
         * @param i
         */
        private void heapify(int i) {
            // 获取左右结点的数组下标 (i + 1) << 1 = 2 * i + 2
            int left = ((i + 1) << 1) - 1;
            int right = (i + 1) << 1;

            // 这是一个临时变量，表示 跟结点、左结点、右结点中最大的值的结点的下标
            int largest = i;

            // 存在左结点，且左结点的值大于根结点的值
            if (left < data.length && data[left] > data[largest]) {
                largest = left;
            }

            // 存在右结点，且右结点的值大于以上比较的较大值
            if (right < data.length && data[right] > data[largest]) {
                largest = right;
            }

            // 左右结点的值都小于根节点，直接return，不做任何操作
            if (i == largest) {
                return;
            }

            // 交换根节点和左右结点中最大的那个值，把根节点的值替换下去
            swap(i, largest);

            // 由于替换后左右子树会被影响，所以要对受影响的子树再进行heapify
            heapify(largest);
        }

        /**
         * 交换元素位置
         * @param i
         * @param j
         */
        private void swap(int i, int j) {
            int tmp = data[i];
            data[i] = data[j];
            data[j] = tmp;
        }

        /**
         * 获取对中的最小的元素，根元素
         * @return
         */
        public int getRoot() {
            return data[0];
        }

        /**
         * 替换根元素，并重新heapify
         * @param root
         */
        public void setRoot(int root) {
            data[0] = root;
            heapify(0);
        }
    }

    public static int[] findLeastK(int[] array, int k) {
        if (array == null || array.length == 0 || k <= 0) {
            return null;
        }
        int[] result = new int[k];
        // 先取出数组的前K个元素来构造最大堆
        for (int i = 0; i < k; i++) {
            result[i] = array[i];
        }
        MaxHeap maxHeap = new MaxHeap(result);
        // 遍历数组其余的元素并替换最大堆的root
        for (int i = k; i < array.length; i++) {
            int max = maxHeap.getRoot();
            if (array[i] < max) {
                maxHeap.setRoot(array[i]);
            }
        }
        return result;
    }

    public static void main(String[] args) {
        int[] array = {4, 5, 1, 6, 2, 7, 3, 8};
        int[] result = findLeastK(array, 4);
        for (int i : result) {
            System.out.print(i + " ");
        }
    }

输出是 4 3 1 2