feat: implement quick select in kth largest element in an array

Author: bujosa

alternative/medium/kth_largest_element_in_an_array.py

@@ -1,6 +1,9 @@
 from typing import List
 import heapq
 
+# Algorithm - Max Heap
+# Time Complexity - O(nlogn)
+# Space Complexity - O(n)
 def findKthLargest(nums: List[int], k: int) -> int:
     max_heap = []
 
@@ -13,9 +16,36 @@ def findKthLargest(nums: List[int], k: int) -> int:
 
     return -heapq.heappop(max_heap)
 
-# Algorithm - Max Heap
-# Time Complexity - O(nlogn)
-# Space Complexity - O(n)
+# Algorithm - Quick Select
+# Time Complexity - O(n)
+# Space Complexity - O(1)
+def findKthLargestV2(nums: List[int], k: int) -> int:
+    k = len(nums) - k
+    left, right = 0, len(nums) - 1
 
+    while left < right:
+        pivot = partition(nums, left, right)
+        if pivot < k:
+            left = pivot + 1
+        elif pivot > k:
+            right = pivot - 1
+        else:
+            break
+
+    return nums[k]
+
+def partition(nums: List[int], left: int, right: int) -> int:
+    pivot = left
+    for i in range(left, right):
+        if nums[i] <= nums[right]:
+            nums[i], nums[pivot] = nums[pivot], nums[i]
+            pivot += 1
+
+    nums[pivot], nums[right] = nums[right], nums[pivot]
+
+    return pivot
+    
 assert findKthLargest([3,2,1,5,6,4], 2) == 5
-assert findKthLargest([3,2,3,1,2,4,5,5,6], 4) == 4
+assert findKthLargestV2([3,2,1,5,6,4], 2) == 5
+assert findKthLargest([3,2,3,1,2,4,5,5,6], 4) == 4
+assert findKthLargestV2([3,2,3,1,2,4,5,5,6], 4) == 4