11.8 Find K'th Largest Element - alternate solution

#1

Hi friends :slight_smile:

I found a solution for this exercise (pages 157-159 in the Python book) which optimizes the partition algorithm. I used a simplified version of the Dutch National Flag partition algorithm (exercise 5.1, pages 41-44 in EPI Python), avoiding swaps for every element that is already in order.

After completing the partition, we return the first index of the right-side interval, consisting of all elements greater than or equal to the pivot value, but the pivot value itself may or may not be at that location. In order to make this work, you have to expand the right-side search window by 1 to include the item at new_pivot_idx, because it may or may not equal the pivot. That increases the search iterations by at least 1, but I correspondingly reduce search iterations by 1 by ending when we have only 1 element left in the search window, which is the solution.

Ending at len(A)==1 also avoids placing the return statement inside the conditional and the IndexError statement which would be required to catch the remaining cases. Finally, as a bonus, it’s 5 significant lines more concise, which is always satisfying! Enjoy <3

import random

def find_kth_largest(k: int, A: List[int]) -> int:
    def dutch_flag_partition_modified(smaller: int, larger_eq: int, pivot: int) -> int:
        while smaller < larger_eq:
            if A[smaller] < pivot:
                smaller += 1
            else:
                larger_eq -= 1 # Increment before access since larger_eq is past end
                A[smaller], A[larger_eq] = A[larger_eq], A[smaller]
        return larger_eq

    # Make "right" index past the end of the array/window (exclusive interval)
    # Solution index is len-k if we're searching for kth largest, or k-1 if kth smallest
    left, right, solution_idx = 0, len(A), len(A) - k

    # Stop when there is only one element left in the window (the solution)
    while left + 1 < right:
        pivot = A[random.randint(left, right-1)]
        new_pivot_idx = dutch_flag_partition_modified(left, right, pivot)
        if new_pivot_idx > solution_idx:
            right = new_pivot_idx
        else:  # new_pivot_idx < solution_idx
            left = new_pivot_idx
    return A[solution_idx]
0 Likes