牛客例题一百遍：寻找第K大

描述

有一个整数数组，请你根据快速排序的思路，找出数组中第$K$大的数。

给定一个整数数组$a$ ,同时给定它的大小n和要找的$K(1\le K\le N)$ ，请返回第$K$大的数(包括重复的元素，不用去重)，保证答案存在。
要求时间复杂度为$O(n)$

示例1

输入：[1,3,5,2,2],5,3
返回值：2

示例2

输入：[10,10,9,9,8,7,5,6,4,3,4,2],12,3
返回值：9
说明：去重后的第3大是8，但本题要求包含重复的元素，不用去重，所以输出9 

解体思路

思路一

堆排序，代码如下：

public int findKth(int[] a, int n, int K) {
        // write code here
        TreeMap<Integer, Integer> heapMap = new TreeMap<>();
        if (n == 0) {
            return -1;
        }
        int count = 0;
        for (int i = 0; i < a.length; i++) {
            if (count <= K) {
                heapMap.put(a[i], heapMap.getOrDefault(a[i], 0) + 1);
                count++;
                continue;
            }
            Map.Entry<Integer, Integer> firstEntry = heapMap.firstEntry();
            if (firstEntry.getKey() < a[i]) {
                if (firstEntry.getValue() == 1) {
                    heapMap.remove(firstEntry.getKey());
                } else {
                    heapMap.put(firstEntry.getKey(), firstEntry.getValue() - 1);
                }
                heapMap.put(a[i], heapMap.getOrDefault(a[i], 0) + 1);
            }
        }
        return heapMap.firstEntry().getKey();
    }

复杂度分析：
时间复杂度：$O(N\log K)$，堆的时间复杂度。
空间复杂度：$O(N)$，需要TreeMap保存数据。

思路二

快排+二分法，代码如下：

      public int findKth(int[] a, int n, int K) {
        // write code here
        return quickSort(a, 0, n - 1, a.length - K);

    }

    private int quickSort(int[] a, int left, int right, int needIndex) {
        int index = partition(a, left, right);

        if (index == needIndex) {
            return a[index];
        }
        if (index > needIndex) {
            return quickSort(a, left, index - 1, needIndex);
        } else {
            return quickSort(a, index + 1, right, needIndex);
        }

    }

    private int partition(int[] a, int left, int right) {
        int tmp = a[left];
        while (left < right) {

            //找到比第一个小的交换过去
            while (left < right && a[right] >= tmp) {
                right--;
            }
            a[left] = a[right];

            //找到比交换过去大的交换过去
            while (left < right && a[left] <= tmp) {
                left++;
            }
            a[right] = a[left];
        }
        a[left] = tmp;
        return left;
    }

复杂度分析：
时间复杂度：$O(N)$，每个节点只遍历一次，所以只需要N复杂度。
空间复杂度：$O(\log N)$，堆栈的空间复杂度。