Wiggle sort

sorts/wiggle_sort.py

@@ -9,18 +9,28 @@
 """
 
 
-def wiggle_sort(nums):
-    """Perform Wiggle Sort."""
-    for i in range(len(nums)):
+def wiggle_sort(nums: list):
+    """
+    Python implementation of wiggle.
+    Example:
+    >>> wiggle_sort([0, 5, 3, 2, 2])
+    [0, 5, 2, 3, 2]
+    >>> wiggle_sort([])
+    []
+    >>> wiggle_sort([-2, -5, -45])
+    [-45, -2, -5]
+    """
+    for i, _ in enumerate(nums):
         if (i % 2 == 1) == (nums[i - 1] > nums[i]):
             nums[i - 1], nums[i] = nums[i], nums[i - 1]
 
+    return nums
+
 
 if __name__ == "__main__":
-    print("Enter the array elements:\n")
+    print("Enter the array elements:")
     array = list(map(int, input().split()))
-    print("The unsorted array is:\n")
-    print(array)
-    wiggle_sort(array)
-    print("Array after Wiggle sort:\n")
+    print("The unsorted array is:")
     print(array)
+    print("Array after Wiggle sort:")
+    print(wiggle_sort(array))