Implement partial_sort_copy.

clibraries · clibraries · commit e1a8b5ec5350 · 2023-07-13T21:04:28.000-06:00
diff --git a/array_alg.h b/array_alg.h
@@ -673,6 +673,15 @@ ALGDEF void NS(partial_sort)(
         void *compare_ctx
         );
 
+ALGDEF T *NS(partial_sort_copy)(
+        const T *first,
+        const T *last,
+        T *out_first,
+        T *out_last,
+        int (*compare)(const T*, const T*, void*),
+        void *compare_ctx
+        );
+
 ALGDEF void NS(nth_element)(
         T *first,
         T *nth,
@@ -2109,6 +2118,42 @@ ALGDEF void NS(partial_sort)(
     NS(sort_heap)(first, middle, compare, compare_ctx);
 }
 
+ALGDEF T *NS(partial_sort_copy)(
+        const T *first,
+        const T *last,
+        T *out_first,
+        T *out_last,
+        int (*compare)(const T*, const T*, void*),
+        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;
+
+    NS(make_heap_n)(out_first, n, 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);
+        }
+        ++first;
+    }
+    NS(sort_heap)(out_first, out_last, compare, compare_ctx);
+    return out_last;
+}
+
 ALGDEF void NS(nth_element)(
         T *first,
         T *nth,
diff --git a/test/tests.c b/test/tests.c
@@ -589,6 +589,25 @@ 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;
+    }
+
+    for (int iteration = 0; iteration < 1000; ++iteration) {
+        intv_random_shuffle(nums, nums + N);
+
+        int top[10];
+        intv_partial_sort_copy(nums, nums + N, top, top + 10, compare_int, NULL);
+
+        for (int i = 0; i < 10; ++i) {
+            assert(top[i] == i);
+        }
+    }
+}
+
 typedef void (*SortFunc)(int*, int*, int (*cmp)(const int*, const int*, void*), void*);
 
 static inline
@@ -811,6 +830,7 @@ int main() {
     printf("-- test_sort_partition -- \n"); test_sort_partition();
     printf("-- test_nth_element --\n"); test_nth_element();
     printf("-- test_partial_sort --\n"); test_partial_sort();
+    printf("-- test_partial_sort_copy --\n"); test_partial_sort_copy();
     printf("-- test_heap_sort --\n"); test_heap_sort();
     printf("-- test_insertion_sort --\n"); test_insertion_sort();
     printf("-- test_stable_sort --\n"); test_stable_sort();
