Reduce Time Complexity and Solve Error.

Author: kadubhumika

data_structures/arrays/kth_largest_element.py

@@ -1,117 +1,46 @@
-"""
-Given an array of integers and an integer k, find the kth largest element in the array.
+import heapq
+from typing import List, Union
 
-https://stackoverflow.com/questions/251781
-"""
 
-
-def partition(arr: list[int], low: int, high: int) -> int:
+def kth_largest_element(arr: List[Union[int, float, str]], k: int) -> Union[int, float, str]:
     """
-    Partitions list based on the pivot element.
-
-    This function rearranges the elements in the input list 'elements' such that
-    all elements greater than or equal to the chosen pivot are on the right side
-    of the pivot, and all elements smaller than the pivot are on the left side.
+    Finds the kth largest element in an array using a heap.
 
     Args:
-        arr: The list to be partitioned
-        low: The lower index of the list
-        high: The higher index of the list
-
-    Returns:
-        int: The index of pivot element after partitioning
-
-        Examples:
-        >>> partition([3, 1, 4, 5, 9, 2, 6, 5, 3, 5], 0, 9)
-        4
-        >>> partition([7, 1, 4, 5, 9, 2, 6, 5, 8], 0, 8)
-        1
-        >>> partition(['apple', 'cherry', 'date', 'banana'], 0, 3)
-        2
-        >>> partition([3.1, 1.2, 5.6, 4.7], 0, 3)
-        1
-    """
-    pivot = arr[high]
-    i = low - 1
-    for j in range(low, high):
-        if arr[j] >= pivot:
-            i += 1
-            arr[i], arr[j] = arr[j], arr[i]
-    arr[i + 1], arr[high] = arr[high], arr[i + 1]
-    return i + 1
-
-
-def kth_largest_element(arr: list[int], position: int) -> int:
-    """
-    Finds the kth largest element in a list.
-    Should deliver similar results to:
-    ```python
-    def kth_largest_element(arr, position):
-        return sorted(arr)[-position]
-    ```
-
-    Args:
-        nums: The list of numbers.
+        arr: List of numbers or strings.
         k: The position of the desired kth largest element.
 
     Returns:
-        int: The kth largest element.
+        The kth largest element.
+
+    Raises:
+        ValueError: If k is invalid or not an integer.
 
     Examples:
-        >>> kth_largest_element([3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5], 3)
+        >>> kth_largest_element([3, 1, 4, 1, 5, 9], 2)
         5
-        >>> kth_largest_element([2, 5, 6, 1, 9, 3, 8, 4, 7, 3, 5], 1)
+        >>> kth_largest_element([9, 7, 5, 3, 1], 1)
         9
-        >>> kth_largest_element([2, 5, 6, 1, 9, 3, 8, 4, 7, 3, 5], -2)
-        Traceback (most recent call last):
-        ...
-        ValueError: Invalid value of 'position'
-        >>> kth_largest_element([9, 1, 3, 6, 7, 9, 8, 4, 2, 4, 9], 110)
-        Traceback (most recent call last):
-        ...
-        ValueError: Invalid value of 'position'
-        >>> kth_largest_element([1, 2, 4, 3, 5, 9, 7, 6, 5, 9, 3], 0)
-        Traceback (most recent call last):
-        ...
-        ValueError: Invalid value of 'position'
-        >>> kth_largest_element(['apple', 'cherry', 'date', 'banana'], 2)
-        'cherry'
-        >>> kth_largest_element([3.1, 1.2, 5.6, 4.7,7.9,5,0], 2)
-        5.6
-        >>> kth_largest_element([-2, -5, -4, -1], 1)
-        -1
+        >>> kth_largest_element([1, 2, 3], 3)
+        1
         >>> kth_largest_element([], 1)
-        -1
-        >>> kth_largest_element([3.1, 1.2, 5.6, 4.7, 7.9, 5, 0], 1.5)
         Traceback (most recent call last):
         ...
-        ValueError: The position should be an integer
-        >>> kth_largest_element((4, 6, 1, 2), 4)
+        ValueError: Invalid value of 'k'
+        >>> kth_largest_element([1, 2, 3], -2)
         Traceback (most recent call last):
         ...
-        TypeError: 'tuple' object does not support item assignment
+        ValueError: Invalid value of 'k'
     """
-    if not arr:
-        return -1
-    if not isinstance(position, int):
+    if not isinstance(k, int):
         raise ValueError("The position should be an integer")
-    if not 1 <= position <= len(arr):
-        raise ValueError("Invalid value of 'position'")
-    low, high = 0, len(arr) - 1
-    while low <= high:
-        if low > len(arr) - 1 or high < 0:
-            return -1
-        pivot_index = partition(arr, low, high)
-        if pivot_index == position - 1:
-            return arr[pivot_index]
-        elif pivot_index > position - 1:
-            high = pivot_index - 1
-        else:
-            low = pivot_index + 1
-    return -1
+    if not 1 <= k <= len(arr):
+        raise ValueError("Invalid value of 'k'")
+
+    # Use heapq.nlargest to directly get k largest, then pick the last
+    return heapq.nlargest(k, arr)[-1]
 
 
 if __name__ == "__main__":
     import doctest
-
     doctest.testmod()