尋找數組中第K大的元素

問題

Find the kth largest element in an unsorted array. Note that it is the kth largest element in the sorted order, not the kth distinct element.

Example 1:

Input: [3,2,1,5,6,4] and k = 2
Output: 5
Example 2:

Input: [3,2,3,1,2,4,5,5,6] and k = 4
Output: 4
Note: 
You may assume k is always valid, 1 ≤ k ≤ array's length. 複製代碼

tag: Medium

分析

這題最簡單的作法是將數組排序，而後直接返回第K大的元素。複雜度爲：O(NlogN)。可是，很明顯，出題者並不想讓咱們這麼作。

若是對數組排序，算法的複雜度起碼是 O(NlogN)。那麼若是咱們不排序，能不能求出第K大元素呢？答案是能夠的，咱們知道快速排序中有一個步驟是 partition。它選擇一個元素做爲樞紐（pivot），將全部小於樞紐的元素放到樞紐的左邊，將全部大於樞紐的元素放到樞紐的右邊。而後返回樞紐在數組中的位置。那麼，關鍵就在這裏了。若是此時返回的樞紐元素在數組中的位置恰好是咱們所要求的位置，問題就能獲得解決了。

咱們先選取一個樞紐元素，從數組的第一個元素開始到最後一個元素結束。將大於樞紐的左邊元素和小於樞紐的右邊元素交換。而後判斷當前樞紐元素的 index 是否爲 n - k，若是是則直接返回這個樞紐元素。若是 index 大於 n - k，則咱們從樞紐的左邊繼續遞歸尋找，若是 index 小於 n - k，則咱們從樞紐的右邊繼續遞歸尋找。

代碼

public class KthLargestElementInArray {

    public int findKthLargest(int[] nums, int k) {
        return findKth(nums, nums.length - k /* 第k大，也就是排序後數組中 index 爲 n - k 的元素*/, 0, nums.length - 1);
    }

    public int findKth(int[] nums, int k, int low, int high) {
        if (low >= high) {
            return nums[low];
        }
        int pivot = partition(nums, low, high); // 樞紐元素的 index
        if (pivot == k) {
            return nums[pivot];
        } else if (pivot < k) { // 樞紐元素的 index 小於 k，繼續從樞紐的右邊部分找
            return findKth(nums, k, pivot + 1, high);
        } else { // 樞紐元素的 index 大於 k，繼續從樞紐的左邊部分找
            return findKth(nums, k, low, pivot - 1);
        }
    }

    int partition(int[] nums, int low, int high) {
        int i = low;
        int j = high + 1;
        int pivot = nums[low];
        while (true) {
            while (less(nums[++i], pivot)) {
                if (i == high) {
                    break;
                }
            }
            while (less(pivot, nums[--j])) {
                if (j == low) {
                    break;
                }
            }
            if (i >= j) {
                break;
            }
            swap(nums, i, j);
        }
        swap(nums, low, j);
        return j;
    }

    boolean less(int i, int j) {
        return i < j;
    }

    void swap(int[] nums, int i, int j) {
        if (i == j) {
            return;
        }
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }
}
複製代碼

複雜度分析

複雜度爲：O(N) + O(N / 2) + O(N / 4) + ... 最終算法複雜度爲 O(N)