LeetCode #215: Kth Largest Element in an Array

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:

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.

Solution:

• Method 1: sort the first k elements and add new elements by binary search

  • Time complexity: O( k*log(k) + (N-k)*log(k) ) = O(N*log(k))
  • Space complexity: O(k)

• Method 2: build a min-heap for the first k elements and add new element one-by-one if the new element is larger than the root of the min-heap, and then heapify the heap.

  • Time complexity: O(k + (N-k)*log(k))
  • Space complexity: O(k)

• Method 3: build a max-heap for the array, then heapify it k-1 time, then the root becomes the kth largest number

  • Time complexity: O(N + (k-1)*log(N))
  • Space complexity: O(N)

class Solution:
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        method = 'bs' # ['bs', 'minheap', 'maxheap']        
        N = len(nums)

        # binary search for inserting the new element
        if method == 'bs':
            tmp = sorted( nums[:k] )
            for i in range(k, len(nums)):
                if nums[i]>=tmp[-1]:
                    tmp.pop(0)
                    tmp.append(nums[i])
                else:
                    if nums[i]>tmp[0]:
                        j = self.binary_search(tmp, 0, k-1, nums[i])
                        tmp.insert(j+1, nums[i])
                        tmp.pop(0)
            return tmp[0]

        # using min-heap
        if method == 'minheap':
            tmp = nums[:k]
            for i in range(k)[::-1]:
                self.heapify_min(tmp, k, i)
            for i in range(k,len(nums)):
                if nums[i] > tmp[0]:
                    tmp[0] = nums[i]
                    self.heapify_min(tmp, k, 0)
            return tmp[0]

        # using max-heap
        if method == 'maxheap':            
            for i in range(N)[::-1]:
                self.heapify_max(nums, N, i)
            for i in range(k-1):
                x = nums.pop(0)
                self.heapify_max(nums, N-i-1, 0)
            return nums[0]

    # binary search
    def binary_search(self, array, start, end, x):
        #print(array, start, end, x)
        if (start+1)==end:
            return start
        else:
            mid = (start+end)//2
            if x < array[mid]:
                return self.binary_search(array, start, mid, x)
            else:
                return self.binary_search(array, mid, end, x)

    # min heap
    def heapify_min(self, array, n, i):
        if i<n:
            imin   = i
            ileft  = 2*i+1
            iright = 2*i+2
            if ileft < n and array[imin] > array[ileft]:
                imin = ileft
            if iright < n and array[imin] > array[iright]:
                imin = iright
            if imin != i:
                array[imin], array[i] = array[i], array[imin]
            self.heapify_min(array, n, ileft)
            self.heapify_min(array, n, iright)        
        return

    # max heap
    def heapify_max(self, array, n, i):
        if i<n:
            imax   = i
            ileft  = 2*i+1
            iright = 2*i+2
            if ileft < n and array[imax] < array[ileft]:
                imax = ileft
            if iright < n and array[imax] < array[iright]:
                imax = iright
            if imax != i:
                array[imax], array[i] = array[i], array[imax]
            self.heapify_max(array, n, ileft)
            self.heapify_max(array, n, iright)        
        return