diff --git a/sorts/stooge_sort.py b/sorts/stooge_sort.py index 767c6a05924f..8947bc3c9911 100644 --- a/sorts/stooge_sort.py +++ b/sorts/stooge_sort.py @@ -1,4 +1,11 @@ -def stooge_sort(arr: list[int]) -> list[int]: +from typing import Protocol + + +class Comparable(Protocol): + def __lt__(self, other: object, /) -> bool: ... + + +def stooge_sort[T: Comparable](arr: list[T]) -> list[T]: """ Examples: >>> stooge_sort([18.1, 0, -7.1, -1, 2, 2]) @@ -6,17 +13,28 @@ def stooge_sort(arr: list[int]) -> list[int]: >>> stooge_sort([]) [] + + >>> stooge_sort(["c", "a", "b"]) + ['a', 'b', 'c'] + + >>> stooge_sort([2.5, -1, 0.0]) + [-1, 0.0, 2.5] + + >>> stooge_sort([1, "a"]) + Traceback (most recent call last): + ... + TypeError: '<' not supported between instances of 'str' and 'int' """ stooge(arr, 0, len(arr) - 1) return arr -def stooge(arr: list[int], i: int, h: int) -> None: +def stooge[T: Comparable](arr: list[T], i: int, h: int) -> None: if i >= h: return # If first element is smaller than the last then swap them - if arr[i] > arr[h]: + if arr[h] < arr[i]: arr[i], arr[h] = arr[h], arr[i] # If there are more than 2 elements in the array