diff --git a/sorts/quick_sort.py b/sorts/quick_sort.py index 374d52e75c81..77b4ecc70892 100644 --- a/sorts/quick_sort.py +++ b/sorts/quick_sort.py @@ -11,9 +11,14 @@ from __future__ import annotations from random import randrange +from typing import Protocol -def quick_sort(collection: list) -> list: +class Comparable(Protocol): + def __lt__(self, other: object, /) -> bool: ... + + +def quick_sort[T: Comparable](collection: list[T]) -> list[T]: """A pure Python implementation of quicksort algorithm. :param collection: a mutable collection of comparable items @@ -26,18 +31,26 @@ def quick_sort(collection: list) -> list: [] >>> quick_sort([-2, 5, 0, -45]) [-45, -2, 0, 5] + >>> quick_sort(["banana", "apple", "cherry"]) + ['apple', 'banana', 'cherry'] + >>> quick_sort([3.14, 1.5, 2.7]) + [1.5, 2.7, 3.14] + >>> quick_sort([1, "two"]) # doctest: +ELLIPSIS + Traceback (most recent call last): + ... + TypeError: '<' not supported between instances of ... """ # Base case: if the collection has 0 or 1 elements, it is already sorted if len(collection) < 2: return collection - # Randomly select a pivot index and remove the pivot element from the collection + # Randomly select a pivot index and remove the pivot element pivot_index = randrange(len(collection)) pivot = collection.pop(pivot_index) - # Partition the remaining elements into two groups: lesser or equal, and greater - lesser = [item for item in collection if item <= pivot] - greater = [item for item in collection if item > pivot] + # Partition the remaining elements using the less-than comparison + lesser = [item for item in collection if item < pivot] + greater = [item for item in collection if not item < pivot] # Recursively sort the lesser and greater groups, and combine with the pivot return [*quick_sort(lesser), pivot, *quick_sort(greater)]