Improved Heap_Sort

sorts/heap_sort.py

@@ -1,58 +1,56 @@
-"""
-This is a pure Python implementation of the heap sort algorithm.
-
-For doctests run following command:
-python -m doctest -v heap_sort.py
-or
-python3 -m doctest -v heap_sort.py
-
-For manual testing run:
-python heap_sort.py
-"""
+def heapify(arr, root, size):
+    """
+    Heapify a subtree rooted at 'root' to maintain the max-heap property.
 
+    Args:
+        arr (list): The input list.
+        root (int): The index of the root element of the subtree.
+        size (int): The size of the heap.
+    """
+    max_index = root
+    left_child = 2 * root + 1
+    right_child = 2 * root + 2
 
-def heapify(unsorted, index, heap_size):
-    largest = index
-    left_index = 2 * index + 1
-    right_index = 2 * index + 2
-    if left_index < heap_size and unsorted[left_index] > unsorted[largest]:
-        largest = left_index
+    if left_child < size and arr[left_child] > arr[max_index]:
+        max_index = left_child
 
-    if right_index < heap_size and unsorted[right_index] > unsorted[largest]:
-        largest = right_index
+    if right_child < size and arr[right_child] > arr[max_index]:
+        max_index = right_child
 
-    if largest != index:
-        unsorted[largest], unsorted[index] = unsorted[index], unsorted[largest]
-        heapify(unsorted, largest, heap_size)
+    if max_index != root:
+        arr[root], arr[max_index] = arr[max_index], arr[root]
+        heapify(arr, max_index, size)
 
 
-def heap_sort(unsorted):
+def heap_sort(arr):
     """
-    Pure implementation of the heap sort algorithm in Python
-    :param collection: some mutable ordered collection with heterogeneous
-    comparable items inside
-    :return: the same collection ordered by ascending
+    Sorts an input list using the heap sort algorithm.
 
-    Examples:
-    >>> heap_sort([0, 5, 3, 2, 2])
-    [0, 2, 2, 3, 5]
+    Args:
+        arr (list): The input list to be sorted.
 
-    >>> heap_sort([])
-    []
-
-    >>> heap_sort([-2, -5, -45])
-    [-45, -5, -2]
+    Returns:
+        list: The sorted list.
     """
-    n = len(unsorted)
+    n = len(arr)
+
+    # Build a max-heap.
     for i in range(n // 2 - 1, -1, -1):
-        heapify(unsorted, i, n)
+        heapify(arr, i, n)
+
+    # Extract elements one by one and build a sorted array.
     for i in range(n - 1, 0, -1):
-        unsorted[0], unsorted[i] = unsorted[i], unsorted[0]
-        heapify(unsorted, 0, i)
-    return unsorted
+        arr[0], arr[i] = arr[i], arr[0]  # Swap root with the last element.
+        heapify(arr, 0, i)  # Call heapify on the reduced heap.
+
+    return arr
 
 
 if __name__ == "__main__":
-    user_input = input("Enter numbers separated by a comma:\n").strip()
-    unsorted = [int(item) for item in user_input.split(",")]
-    print(heap_sort(unsorted))
+    user_input = input("Enter numbers separated by commas: ").strip()
+    try:
+        unsorted = [int(item) for item in user_input.split(",")]
+    except ValueError:
+        print("Invalid input. Please enter a list of integers separated by commas.")
+        exit()
+    print(f"Sorted list: {heap_sort(unsorted)}")