剑指Offer-30. 最小的K个数

题目描述

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

思路

用堆排序，O(nlogk)

import java.util.*;

public class Solution {
    public ArrayList<Integer> GetLeastNumbers_Solution(int[] input, int k) {

        int len = input.length;

        ArrayList<Integer> res = new ArrayList<>();

        if (input == null || k <= 0 || k > len)
            return res;

        // 取前 k 个数字
        int[] kArray = Arrays.copyOfRange(input, 0, k);
        // 对前 k 个数字进行堆排序
        heapSort(kArray);

        // 从第 k+1 个数字开始比较
        for (int i = k; i < len; i++) {
            // 如果第 k+1 个数比大根堆的根结点小，就取代根结点，然后重新堆排序
            if (input[i] < kArray[k - 1]) {
                kArray[k - 1] = input[i];
                heapSort(kArray);
            }
        }

        for (int i : kArray) {
            res.add(i);
        }

        return res;
    }

    public void adjustHeap(int[] arr, int i, int length) {

        int tmp = arr[i];

        for (int k = 2*i + 1; k < length; k = 2*k + 1) {
            if (k + 1 < length && arr[k + 1] > arr[k])
                k++;
            if (arr[k] > tmp) {
                arr[i] = arr[k];
                i = k;
            } else {
                break;
            }
        }

        arr[i] = tmp;
    }

    public void heapSort(int[] arr) {

        int len = arr.length;

        for (int i = len/2 - 1; i >= 0; i--) {
            adjustHeap(arr, i, len);
        }

        for (int i = len - 1; i > 0; i--) {
            swap(arr, 0, i);
            adjustHeap(arr, 0, i);
        }
    }

    public void swap(int[]arr, int i, int j) {
        int tmp = arr[i];
        arr[i] = arr[j];
        arr[j] = tmp;
    }
}