Problem statement

https://binarysearch.com/problems/Kth-Smallest-Element/

Solution

Variation of Leetcode 0215. Kth Largest Element in an Array, but we need to find smaller here, and index starts with 0.

Complexity

It is O(n) in average

Code

class Solution:
    def solve(self, nums, k):
        pivot = random.choice(nums)
        left =  [x for x in nums if x < pivot]
        mid  =  [x for x in nums if x == pivot]
        right = [x for x in nums if x > pivot]
        
        L, M = len(left), len(mid)
        
        if k < L:
            return self.solve(left, k)
        elif k >= L + M:
            return self.solve(right, k - L - M)
        else:
            return mid[0]