Use with_bit_floor in radix sort

cp-algorithms · Nov 14, 2024 · dc7b822 · dc7b822
1 parent 52d8541
commit dc7b822
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 8 deletions.
diff --git a/cp-algo/util/bit.hpp b/cp-algo/util/bit.hpp
@@ -15,5 +15,15 @@ namespace cp_algo {
     size_t kth_set_bit(uint64_t x, size_t k) {
         return std::countr_zero(_pdep_u64(1ULL << k, x));
     }
+    template<int fl = 0>
+    void with_bit_floor(size_t n, auto &&callback) {
+        if constexpr (fl >= 63) {
+            return;
+        } else if (n >> (fl + 1)) {
+            with_bit_floor<fl + 1>(n, callback);
+        } else {
+            callback.template operator()<1ULL << fl>();
+        }
+    }
 }
 #endif // CP_ALGO_UTIL_BIT_HPP
diff --git a/cp-algo/util/sort.hpp b/cp-algo/util/sort.hpp
@@ -1,12 +1,14 @@
 #ifndef CP_ALGO_UTIL_SORT_HPP
 #define CP_ALGO_UTIL_SORT_HPP
+#include "bit.hpp"
 #include <algorithm>
 #include <numeric>
 #include <ranges>
 #include <vector>
 namespace cp_algo {
-    void count_sort(auto &a, size_t maxc, auto &&proj = std::identity{}) {
-        std::vector<int> cnt(maxc);
+    template<size_t maxc>
+    void count_sort(auto &a, auto &&proj = std::identity{}) {
+        std::array<int, maxc> cnt = {};
         for(auto &x: a) {
             cnt[proj(x)]++;
         }
@@ -22,13 +24,15 @@ namespace cp_algo {
         if(empty(a)) {
             return;
         }
-        int base = std::bit_floor(size(a));
         auto mx = std::ranges::max(a);
-        for(int64_t i = 1; i <= mx; i *= base) {
-            count_sort(a, base, [&](auto x) {
-                return x / i % base;
-            });
-        }
+        with_bit_floor(size(a), [&]<size_t floor>() {
+            constexpr int base = std::min<size_t>(floor, 1 << 16);
+            for(int64_t i = 1; i <= mx; i *= base) {
+                count_sort<base>(a, [&](auto x) {
+                    return x / i % base;
+                });
+            }
+        });
     }
 }
 #endif // CP_ALGO_UTIL_SORT_HPP