Fork me on GitHub

Kth Largest Element in an Array

Posted on

Description

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:

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

Example 2:

1
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.

Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
class Solution(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        return self.find_inner_k_largest(nums, 0, len(nums)-1, k)
        
    def find_inner_k_largest(self, nums, start, end, k):
        lens = end - start + 1
        if k > lens:
            return 
        if lens == 1:
            return nums[start]
        
        middle_index = self.partition(nums, start, end)
        
        if lens - (middle_index-start) == k:
            return nums[middle_index]
        
        if lens - (middle_index-start) < k:
            return self.find_inner_k_largest(nums, start, middle_index-1, k-(end-middle_index+1))
        
        
        return self.find_inner_k_largest(nums, middle_index+1, end, k)
        
        
    ## The single direction partition algorithm is easier than the bi-diretion case
    def partition(self, nums, start, end):
        segment = nums[end]
        lens = end - start + 1
        
        smaller = start - 1
        larger = start 
        while larger <= end:
            if nums[larger] <= segment:
                smaller += 1
                temp = nums[smaller]
                nums[smaller] = nums[larger]
                nums[larger] = temp
            larger += 1
            
        return smaller