LintCode 606. 第K大的元素 II

原题

解

第一步，万年不变的查错。如果给的array是null或空，直接return 0

    public int kthLargestElement2(int[] nums, int k) {
        if (nums == null || nums.length == 0) {
            return -1;
        }
        ...
    }

很简单，还是PriorityQueue，直接上code了。

    /*
     * @param nums: an integer unsorted array
     * @param k: an integer from 1 to n
     * @return: the kth largest element
     */
    public int kthLargestElement2(int[] nums, int k) {
        if (nums == null || nums.length == 0) {
            return -1;
        }
        
        Queue<Integer> pq = new PriorityQueue<Integer>();
        for(int i = 0; i < nums.length; i++) {
            pq.add(nums[i]);
            if (pq.size() > k) {
                pq.poll();
            }
        }
        return pq.poll();
    }
}

解2

这个题其实也可以用QuickSelect来解。原理就是跟QuickSort一样，选一个Pivot Point，然后分开，比它大的去一边，比它小的去一边。关键是这之后，判断一下第k大的在哪边，然后去那边再排序，不管另外一边了。这样的话，理想状态下，每次都能扔掉一半。

    private int quickSelect(int[] nums, int start, int end, int k) {
        int left = start;
        int right = end;
        int pivot = nums[start + (end - start) / 2];
        
        while (left <= right) {
            while (left <= right && nums[left] > pivot) {
                left++;
            }
            
            while (left <= right && nums[right] < pivot) {
                right--;
            }
            
            if (left <= right) {
                swap(nums, left, right);
                left++;
                right--;
            }
        }
        
        if (start + k - 1 <= right) {
            return quickSelect(nums, start, right, k);
        }
        
        if (start + k - 1 >= left) {
            return quickSelect(nums, left, end, k - (left - start));
        }
        
        return nums[right + 1];
    }
    
    private void swap(int[] nums, int a, int b) {
        int temp = nums[a];
        nums[a] = nums[b];
        nums[b] = temp;
    }

这里几个关键的地方跟排序不一样。首先因为是要找第k大的，即降序，所以左边是大的，右边是小的，碰到一样的不继续前进，这样到最后左右两边同时停在PivotPoint上。进行一次无用的交换后，左右各前进一次，将array划分为3段，中间只有一个数字，就是pivot。然后就是判断第k个在哪。从最左边开始，第k个如果超过pivot，就右边，如果小于pivot，就左边。如果就是，那就找到了。这里不能直接认为就是第k个，因为除了pivot，其他的数都没有排序。
完整的code

public class Solution {
    /*
     * @param nums: an integer unsorted array
     * @param k: an integer from 1 to n
     * @return: the kth largest element
     */
    public int kthLargestElement2(int[] nums, int k) {
        if (nums == null || nums.length == 0) {
            return -1;
        }
        
        return quickSelect(nums, 0, nums.length - 1, k);
    }
    
    private int quickSelect(int[] nums, int start, int end, int k) {
        int left = start;
        int right = end;
        int pivot = nums[start + (end - start) / 2];
        
        while (left <= right) {
            while (left <= right && nums[left] > pivot) {
                left++;
            }
            
            while (left <= right && nums[right] < pivot) {
                right--;
            }
            
            if (left <= right) {
                swap(nums, left, right);
                left++;
                right--;
            }
        }
        
        if (start + k - 1 <= right) {
            return quickSelect(nums, start, right, k);
        }
        
        if (start + k - 1 >= left) {
            return quickSelect(nums, left, end, k - (left - start));
        }
        
        return nums[right + 1];
    }
    
    private void swap(int[] nums, int a, int b) {
        int temp = nums[a];
        nums[a] = nums[b];
        nums[b] = temp;
    }
}

分析

时间复杂度

第一个解法用PriorityQueue，n个数，queue的大小为k，所以是O(nlogk)。第二个解法，理想状态是O(n)，但是最坏的情况是O(n²)，例如当每一次的pivot都正好是最小的数，而k=1的时候。所以相比之下，因为题目表明n比k无限大，所以O(nlogk)比O(n)&O(n²)要好。

空间复杂度

第一个解法用到了size为k的queue，所以是O(k)。第二个是in place，所以是O(1)。

如果面试官没有说N远大于k，那么就直接QuickSelect，因为期望的复杂度是O(n)。如果他followup说了，那就用PriorityQueue。或者反过来，先PriorityQueue，然后优化用QuickSelect。不过QuickSelect也不是一般能当场想出来的算法。