Fix introsort heap fallback ignoring range boundaries (#15435)

* fix: restrict introsort heap fallback to the active range

* test: cover introsort heap fallback range boundaries

* fix: sort introsort heap fallback ranges without copying

* test: cover introsort heap sort range boundaries and default end

Author: Miladkhoshdel

sorts/intro_sort.py

@@ -44,29 +44,40 @@ def insertion_sort[T: Comparable](
 
 
 def heapify[T: Comparable](
-    array: list[T], index: int, heap_size: int
+    array: list[T], index: int, heap_size: int, start: int = 0
 ) -> None:  # Max Heap
     """
+    Restore the max heap rooted at an index relative to start.
+    heap_size is the number of elements in the heap beginning at start.
+
     >>> array = [4, 2, 6, 8, 1, 7, 8, 22, 14, 56, 27, 79, 23, 45, 14, 12]
     >>> heapify(array, len(array) // 2, len(array))
     """
     largest = index
     left_index = 2 * index + 1  # Left Node
     right_index = 2 * index + 2  # Right Node
 
-    if left_index < heap_size and array[largest] < array[left_index]:
+    if left_index < heap_size and array[start + largest] < array[start + left_index]:
         largest = left_index
 
-    if right_index < heap_size and array[largest] < array[right_index]:
+    if right_index < heap_size and array[start + largest] < array[start + right_index]:
         largest = right_index
 
     if largest != index:
-        array[index], array[largest] = array[largest], array[index]
-        heapify(array, largest, heap_size)
+        array[start + index], array[start + largest] = (
+            array[start + largest],
+            array[start + index],
+        )
+        heapify(array, largest, heap_size, start)
 
 
-def heap_sort[T: Comparable](array: list[T]) -> list[T]:
+def heap_sort[T: Comparable](
+    array: list[T], start: int = 0, end: int | None = None
+) -> list[T]:
     """
+    Sort the half-open range [start:end] in place and return the original list.
+    If end is omitted, sort through the end of the list.
+
     >>> heap_sort([4, 2, 6, 8, 1, 7, 8, 22, 14, 56, 27, 79, 23, 45, 14, 12])
     [1, 2, 4, 6, 7, 8, 8, 12, 14, 14, 22, 23, 27, 45, 56, 79]
     >>> heap_sort([-2, -11, 0, 0, 0, 87, 45, -69, 78, 12, 10, 103, 89, 52])
@@ -76,14 +87,16 @@ def heap_sort[T: Comparable](array: list[T]) -> list[T]:
     >>> heap_sort([6.2, -45.54, 8465.20, 758.56, -457.0, 0, 1, 2.879, 1.7, 11.7])
     [-457.0, -45.54, 0, 1, 1.7, 2.879, 6.2, 11.7, 758.56, 8465.2]
     """
-    n = len(array)
+    if end is None:
+        end = len(array)
+    n = end - start
 
-    for i in range(n // 2, -1, -1):
-        heapify(array, i, n)
+    for i in range(n // 2 - 1, -1, -1):
+        heapify(array, i, n, start)
 
     for i in range(n - 1, 0, -1):
-        array[i], array[0] = array[0], array[i]
-        heapify(array, 0, i)
+        array[start + i], array[start] = array[start], array[start + i]
+        heapify(array, 0, i, start)
 
     return array
 
@@ -187,7 +200,7 @@ def intro_sort[T: Comparable](
     """
     while end - start > size_threshold:
         if max_depth == 0:
-            return heap_sort(array)
+            return heap_sort(array, start, end)
         max_depth -= 1
         pivot = median_of_3(array, start, start + ((end - start) // 2) + 1, end - 1)
         p = partition(array, start, end, pivot)

tests/test_sorts.py

@@ -34,6 +34,8 @@
 from sorts.gnome_sort import gnome_sort
 from sorts.heap_sort import heap_sort
 from sorts.insertion_sort import insertion_sort
+from sorts.intro_sort import heap_sort as intro_heap_sort
+from sorts.intro_sort import intro_sort as intro_sort_range
 from sorts.intro_sort import sort as intro_sort
 from sorts.iterative_merge_sort import iter_merge_sort
 from sorts.merge_insertion_sort import merge_insertion_sort
@@ -66,6 +68,33 @@ def test_heap_sort() -> None:
     assert heap_sort([5, 4, 3, 2, 1]) == [1, 2, 3, 4, 5]
 
 
+@pytest.mark.parametrize(
+    ("start", "end"),
+    [(start, end) for start in range(7) for end in [None, *range(start, 7)]],
+)
+def test_intro_heap_sort_range(start: int, end: int | None) -> None:
+    collection = [100, 4, 1, 3, 1, -100]
+    expected = collection[:start] + sorted(collection[start:end])
+    if end is not None:
+        expected += collection[end:]
+
+    result = intro_heap_sort(collection, start, end)
+
+    assert result is collection
+    assert collection == expected
+
+
+@pytest.mark.parametrize("max_depth", [0, 1])
+def test_intro_sort_heap_fallback_preserves_surrounding_items(max_depth: int) -> None:
+    collection = [100, *range(40, 0, -1), -100]
+    expected = [100, *range(1, 41), -100]
+
+    result = intro_sort_range(collection, 1, 41, 16, max_depth)
+
+    assert result is collection
+    assert collection == expected
+
+
 SORTS = (
     binary_insertion_sort,
     bubble_sort_iterative,