diff --git a/sorts/wiggle_sort.py b/sorts/wiggle_sort.py index cb3c3a47ab21..fd5be7d9926f 100644 --- a/sorts/wiggle_sort.py +++ b/sorts/wiggle_sort.py @@ -8,8 +8,18 @@ one possible Wiggle Sorted answer is [3, 5, 1, 6, 2, 4]. """ +from __future__ import annotations -def wiggle_sort(nums: list) -> list: +from typing import Any, Protocol + + +class Comparable(Protocol): + def __lt__(self, other: Any, /) -> bool: ... + + def __gt__(self, other: Any, /) -> bool: ... + + +def wiggle_sort[T: Comparable](nums: list[T]) -> list[T]: """ Python implementation of wiggle sort. Reorders an array such that nums[0] <= nums[1] >= nums[2] <= nums[3]... @@ -23,6 +33,12 @@ def wiggle_sort(nums: list) -> list: [-5, -2, -45] >>> wiggle_sort([-2.1, -5.68, -45.11]) [-5.68, -2.1, -45.11] + >>> wiggle_sort(["d", "a", "c", "b"]) + ['a', 'd', 'b', 'c'] + >>> wiggle_sort([1, "a"]) + Traceback (most recent call last): + ... + TypeError: '>' not supported between instances of 'int' and 'str' """ for i in range(1, len(nums)): if (i % 2 == 1 and nums[i - 1] > nums[i]) or (