Leetcode 215: Kth Largest Element in an Array

https://leetcode.com/problems/kth-largest-element-in-an-array/

Code 1 (Using partition idea from quicksort with O(N) / O(N^2) time complexity on average/ worst case and O(1) space. Please see reference on complexity analysis.)

class Solution(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        if not nums:
            return -1
        
        left, right = 0, len(nums)-1
        while True:
            pos = self.partition(nums, left, right)
            if pos == k-1:
                return nums[pos]
            elif pos < k-1:
                left = pos+1
            else:
                right = pos-1    

    def partition(self, nums, left, right):
        # always pick nums[left] as pivot
        pivot = nums[left]
        p1, p2 = left+1, right
        while p1 <= p2:
            if nums[p1]<pivot and nums[p2]>pivot:
                nums[p1], nums[p2] = nums[p2], nums[p1]
                p1, p2 = p1+1, p2-1
            if nums[p1] >= pivot:
                p1 += 1
            if nums[p2] <= pivot:
                p2 -= 1
        
        nums[left], nums[p2] = nums[p2], nums[left]
        return p2

s = Solution()
print s.findKthLargest([7,6,5,4,3,2,1], 5)

The function `partition` puts the numbers smaller than `nums[left]` to its left and then returns the new index of `nums[left]`. The returned index is actually telling us how small `nums[left]` ranks in nums. 

 

Code 2 (Using a heap of size k with O(NlogK) time complexity and O(K) space complexity. ):

from heapq import *
class Solution(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        if not nums:
            return -1
        
        h = []
        for i in xrange(len(nums)):
            if len(h) < k:
                heappush(h, nums[i])
            else:
                if h[0] < nums[i]:
                    heappop(h)
                    heappush(h, nums[i])

        return h[0]

 

Idea

In Python heapq, heap is a min-heap. Therefore, h[0] is always the smallest element in the heap and heappop always pops the smallest element. heappop and heappush takes O(logK) time complexity therefore the total time bound is O(NlogK).

 

Reference:

https://leetcode.com/discuss/38336/solutions-partition-priority_queue-multiset-respectively