Kth Smallest by Partitioning
Problem
Given an unsorted list of integers and a number k, return the k-th smallest value, counting from 1 and counting repeated values separately. Sorting takes O(n log n) time; aim for O(n) expected time by reusing the partition step from quicksort. Leave the caller's list unchanged.
Examples
Input: nums = [7, 2, 9, 4, 4, 1], k = 3
Output: 4
Why: in sorted order the list reads 1, 2, 4, 4, 7, 9Input: nums = [3, 3, 3, 3], k = 2
Output: 3
Why: every value equals the pivot, which must not slow the search downInput: nums = [5], k = 1
Output: 5
Why: edge case, a single valueHints
0 / 3
After quicksort partitions around a pivot, the pivot sits in its final sorted position. Do you then need to sort both sides?
Only the side that contains index k - 1 matters, so the other side can be dropped. Choosing the pivot at random makes a very lopsided split unlikely.
Partition the live range into three parts: smaller than the pivot, equal to it and larger. If index k - 1 falls in the middle part the pivot is the answer; otherwise narrow the range to the part that contains it and repeat.
Solution
Quickselect is quicksort that continues into one side only. A three-way partition gathers every copy of the pivot into a middle band, so the pivot is the answer whenever index k - 1 lands in that band, and a list full of repeated values cannot push the search into quadratic time. A random pivot shrinks the range by a constant fraction on average, so the work adds up to a geometric series and is O(n) expected; the unlikely worst case is O(n²). Space is O(n) for the copy that keeps the caller's list unchanged.
import random
def kth_smallest(nums, k):
a, lo, hi, want = list(nums), 0, len(nums) - 1, k - 1 # want is an index
while True:
pivot = a[random.randint(lo, hi)]
lt, i, gt = lo, lo, hi # a[lo:lt] < pivot, a[gt+1:hi+1] > pivot
while i <= gt:
if a[i] < pivot:
a[lt], a[i] = a[i], a[lt]
lt, i = lt + 1, i + 1
elif a[i] > pivot:
a[gt], a[i] = a[i], a[gt]
gt -= 1
else:
i += 1
if want < lt:
hi = lt - 1 # the answer is among the smaller values
elif want > gt:
lo = gt + 1 # the answer is among the larger values
else:
return pivot
print(kth_smallest([7, 2, 9, 4, 4, 1], 3)) # -> 4
print(kth_smallest([3, 3, 3, 3], 2)) # -> 3
print(kth_smallest([5], 1)) # -> 5Stuck on the idea rather than the code? Quick Sort covers it.