面试题40：最小的 k 个数

import java.util.Arrays;

/**
 * Created by clearbug on 2018/2/26.
 *
 * 面试题40：最小的 k 个数
 *
 * 注意：因为前两天在陌陌面试时被问到的问题是最大的 k 个数，所以这里的代码也是求解最大的 k 个数的，最小的 k 个数原理是一样的。
 */
public class Solution {

    public static void main(String[] args) throws InterruptedException {
        Solution s = new Solution();

        int[] arr = {1, 2, 3, 11, 8, 10, 5, 4, 22, 66, 23, 20};
        System.out.println(Arrays.toString(s.topK1(arr, 5)));
        int[] arr2 = {3, 2, 1, 11, 8, 10, 5, 4, 22, 66, 23, 20};
        System.out.println(Arrays.toString(s.topK2(arr2, 5)));
    }

    /**
     * 方法一就是利用快速排序的思想
     *
     * 时间复杂度：O(n)
     *
     * @param arr
     * @param k
     * @return
     */
    public int[] topK1(int[] arr, int k) {
        int topK = arr.length - k;
        int start = 0;
        int end = arr.length - 1;
        int index = partition(arr, start, end);

        while (index != topK) {
            if (index > topK) {
                end = index - 1;
                index = partition(arr, start, end);
            } else {
                start = index + 1;
                index = partition(arr, start, end);
            }
        }

        return subArr(arr, topK, arr.length);
    }

    /**
     * 方法二就是利用堆排序来解决了
     *
     * 时间复杂度：O(nlogk)
     *
     * @param arr
     * @param k
     * @return
     */
    public int[] topK2(int[] arr, int k) {
        int[] heap = new int[k];
        int heapCount = 0;
        for (int i = 0; i < arr.length; i++) {
            if (heapCount < k) {
                heap[heapCount++] = arr[i];
            }
            if (heapCount == k) { // 初始化 heap 为小顶堆
                heapify(heap);
                heapCount++;
                continue;
            }
            if (heapCount > k) {
                if (arr[i] > heap[0]) {
                    heap[0] = arr[i];
                    // 调整 heap 为小顶堆
                    adjust(heap, 0);
                }
            }
        }
        return heap;
    }

    private void adjust(int[] arr, int i) {
        int leftChildIndex = 2 * i + 1;
        int rightChildIndex = 2 * i + 2;
        int smallestIndex = i;
        if (leftChildIndex < arr.length && arr[leftChildIndex] < arr[smallestIndex]) {
            smallestIndex = leftChildIndex;
        }
        if (rightChildIndex < arr.length && arr[rightChildIndex] < arr[smallestIndex]) {
            smallestIndex = rightChildIndex;
        }
        if (smallestIndex != i) {
            int temp = arr[i];
            arr[i] = arr[smallestIndex];
            arr[smallestIndex] = temp;
            adjust(arr, smallestIndex);
        }
    }

    private void heapify(int[] arr) {
        for (int i = arr.length / 2 - 1; i >= 0; i--) {
            int leftChildIndex = 2 * i + 1;
            int rightChildIndex = 2 * i + 2;
            int smallestIndex = i;
            if (leftChildIndex < arr.length && arr[leftChildIndex] < arr[smallestIndex]) {
                smallestIndex = leftChildIndex;
            }
            if (rightChildIndex < arr.length && arr[rightChildIndex] < arr[smallestIndex]) {
                smallestIndex = rightChildIndex;
            }
            if (smallestIndex != i) {
                int temp = arr[i];
                arr[i] = arr[smallestIndex];
                arr[smallestIndex] = temp;
            }
        }
    }

    private int[] subArr(int[] arr, int start, int end) {
        int[] res = new int[end - start];
        for (int i = start; i < end; i++) {
            res[i - start] = arr[i];
        }
        return res;
    }

    private int partition(int[] arr, int start, int end) {
        int privot = arr[start];
        while (start < end) {
            while (arr[end] >= privot && end > start) {
                end--;
            }
            arr[start] = arr[end];
            while (arr[start] <= privot && end > start) {
                start++;
            }
            arr[end] = arr[start];
        }
        arr[start] = privot;
        return start;
    }
}