Improve partial_sort_copy

clibraries · clibraries · commit d88104c1c625 · 2023-07-13T21:19:43.000-06:00
diff --git a/array_alg.h b/array_alg.h
@@ -2127,31 +2127,36 @@ ALGDEF T *NS(partial_sort_copy)(
         void *compare_ctx
         )
 {
-    size_t n1 = last - first;
-    size_t n2 = out_last - out_first;
-
-    size_t n = (n2 < n1) ? n2 : n1;
-    out_last = out_first + n;
-
-    if (n == 0) { return out_last; }
-
-    NS(copy_n)(first, n, out_first);
-    first += n;
+    if (first == last || out_first == out_last) return out_first;
+    T *out = out_first;
+    *out = *first;
+    ++out;
+    ++first;
 
-    NS(make_heap_n)(out_first, n, compare, compare_ctx);
+    while (out != out_last)
+    {
+        if (first == last) {
+             NS(sort_heap)(out_first, out, compare, compare_ctx);
+             return out;
+        }
+        *out = *first;
+        ++first;
+        ++out;
+        NS(push_heap)(out_first, out, compare, compare_ctx);
+    }
 
     while (first != last)
     {
         if (compare(first, out_first, compare_ctx) < 0)
         {
-            NS(pop_heap_n)(out_first, n, compare, compare_ctx);
-            *(out_last - 1) = *first;
-            NS(push_heap_n)(out_first, n, compare, compare_ctx);
+            NS(pop_heap)(out_first, out, compare, compare_ctx);
+            *(out - 1) = *first;
+            NS(push_heap)(out_first, out, compare, compare_ctx);
         }
         ++first;
     }
-    NS(sort_heap)(out_first, out_last, compare, compare_ctx);
-    return out_last;
+    NS(sort_heap)(out_first, out, compare, compare_ctx);
+    return out;
 }
 
 ALGDEF void NS(nth_element)(
diff --git a/test/tests.c b/test/tests.c
@@ -590,19 +590,40 @@ void test_partial_sort() {
 }
 
 void test_partial_sort_copy() {
-    enum { N = 100 };
-    int nums[N];
-    for (int i = 0; i < N; ++i) {
-        nums[i] = i;
+    {
+        enum { N = 100 };
+        int nums[N];
+        for (int i = 0; i < N; ++i) {
+            nums[i] = i;
+        }
+
+        enum { M = 10 };
+        for (int iteration = 0; iteration < 1000; ++iteration) {
+            intv_random_shuffle(nums, nums + N);
+
+            int top[M];
+            intv_partial_sort_copy(nums, nums + N, top, top + M, compare_int, NULL);
+
+            for (int i = 0; i < M; ++i) {
+                assert(top[i] == i);
+            }
+        }
     }
+    {
+        enum { N = 10 };
+        int nums[N];
+        for (int i = 0; i < N; ++i) {
+            nums[i] = i;
+        }
+        intv_reverse(nums, nums + N);
 
-    for (int iteration = 0; iteration < 1000; ++iteration) {
-        intv_random_shuffle(nums, nums + N);
+        enum { M = 100 };
 
-        int top[10];
-        intv_partial_sort_copy(nums, nums + N, top, top + 10, compare_int, NULL);
+        int top[M];
+        int *out = intv_partial_sort_copy(nums, nums + N, top, top + M, compare_int, NULL);
+        assert(out == top + N);
 
-        for (int i = 0; i < 10; ++i) {
+        for (int i = 0; i < N; ++i) {
             assert(top[i] == i);
         }
     }
