Wavelet nit (#146)

* nit to wavelet matrix

* [auto-verifier] verify commit adaa1d2

* update docs

* [auto-verifier] verify commit 60ca7c0

* update docs for this

* [auto-verifier] verify commit e79d2b8

* fix

* fix

* [auto-verifier] verify commit 5e31745

* minor golf

---------

Co-authored-by: GitHub <noreply@github.com>

Author: lrvideckis

.verify-helper/timestamps.remote.json

@@ -24,7 +24,7 @@
 "tests/library_checker_aizu_tests/data_structures/implicit_seg_tree.test.cpp": "2024-11-17 14:04:03 -0600",
 "tests/library_checker_aizu_tests/data_structures/kruskal_tree_aizu.test.cpp": "2025-02-10 23:30:47 -0700",
 "tests/library_checker_aizu_tests/data_structures/kth_smallest_pst.test.cpp": "2024-12-05 10:41:42 -0600",
-"tests/library_checker_aizu_tests/data_structures/kth_smallest_wavelet_matrix.test.cpp": "2024-11-19 08:31:51 -0600",
+"tests/library_checker_aizu_tests/data_structures/kth_smallest_wavelet_matrix.test.cpp": "2025-07-10 11:33:52 -0600",
 "tests/library_checker_aizu_tests/data_structures/lazy_segment_tree.test.cpp": "2024-12-14 19:50:29 -0600",
 "tests/library_checker_aizu_tests/data_structures/lazy_segment_tree_constructor.test.cpp": "2024-12-14 19:50:29 -0600",
 "tests/library_checker_aizu_tests/data_structures/lazy_segment_tree_inc.test.cpp": "2024-12-05 10:41:42 -0600",
@@ -74,7 +74,7 @@
 "tests/library_checker_aizu_tests/handmade_tests/functional_graph.test.cpp": "2025-07-08 19:28:25 -0600",
 "tests/library_checker_aizu_tests/handmade_tests/lca_ladder_forest.test.cpp": "2025-02-10 23:30:47 -0700",
 "tests/library_checker_aizu_tests/handmade_tests/manacher.test.cpp": "2024-12-14 19:50:29 -0600",
-"tests/library_checker_aizu_tests/handmade_tests/merge_st_and_wavelet.test.cpp": "2025-02-11 13:53:30 -0700",
+"tests/library_checker_aizu_tests/handmade_tests/merge_st_and_wavelet.test.cpp": "2025-07-10 11:33:52 -0600",
 "tests/library_checker_aizu_tests/handmade_tests/mobius.test.cpp": "2025-02-10 14:50:36 -0700",
 "tests/library_checker_aizu_tests/handmade_tests/mod_int.test.cpp": "2024-12-14 19:50:29 -0600",
 "tests/library_checker_aizu_tests/handmade_tests/n_choose_k.test.cpp": "2025-01-15 00:22:31 -0700",

library/data_structures/seg_tree_uncommon/wavelet_bit_vec.hpp

@@ -1,20 +1,22 @@
 #pragma once
-#define ull uint64_t
+//! @code
+//!   vector<bool> a(n);
+//!   bit_vec bv(a);
+//!   bv.cnt(r); // a[0] + a[1] + ... + a[r - 1]
+//! @endcode
+//! @time O(n + q)
+//! @space O(n / 64)
+using ull = uint64_t;
 struct bit_vec {
   vector<pair<ull, int>> b;
-  //! @time O(n)
-  //! @space O(n / 64)
   bit_vec(const vector<bool>& a): b(sz(a) / 64 + 1) {
     rep(i, 0, sz(a)) b[i >> 6].first |= ull(a[i])
       << (i & 63);
     rep(i, 0, sz(b) - 1) b[i + 1].second =
       popcount(b[i].first) + b[i].second;
   }
-  //! @returns !a[0] + !a[1] + ... + !a[r - 1]
-  //! @time O(1)
-  //! @space O(1)
-  int cnt0(int r) {
+  int cnt(int r) {
     auto [x, y] = b[r >> 6];
-    return r - y - popcount(x & ((1ULL << (r & 63)) - 1));
+    return y + popcount(x & ((1ULL << (r & 63)) - 1));
   }
 };

library/data_structures/seg_tree_uncommon/wavelet_count_less.hpp

@@ -5,11 +5,11 @@
 int count(int l, int r, ull ub) {
   int res = 0;
   for (int h = sz(bv); h--;) {
-    int l0 = bv[h].cnt0(l), r0 = bv[h].cnt0(r);
+    int l0 = bv[h].cnt(l), r0 = bv[h].cnt(r);
     if ((~ub >> h) & 1) l = l0, r = r0;
     else
-      res += r0 - l0, l += bv[h].cnt0(n) - l0,
-        r += bv[h].cnt0(n) - r0;
+      res += r0 - l0, l += bv[h].cnt(n) - l0,
+        r += bv[h].cnt(n) - r0;
   }
   return res;
 }

library/data_structures/seg_tree_uncommon/wavelet_matrix.hpp

@@ -1,35 +1,34 @@
 #pragma once
+//! @code
+//!   vector<ull> a(n);
+//!   wavelet_matrix wm(a, 1e9); // requires a[i] <= 1e9
+//!   wm.kth(l, r, k); //(k+1)th smallest number in [l,r)
+//!   wm.kth(l, r, 0); //min in [l,r)
+//! @endcode
+//! @time O(n * log(max_val) + q * log(max_val))
+//! @space O(n * log(max_val) / 64)
 #include "wavelet_bit_vec.hpp"
 struct wavelet_matrix {
   int n;
   vector<bit_vec> bv;
-  //! Requires a[i] <= max_val
-  //! @time O(n * log(max_val))
-  //! @space O(n * log(max_val) / 64)
   wavelet_matrix(vector<ull> a, ull max_val):
     n(sz(a)), bv(bit_width(max_val), {{}}) {
-    vector<ull> nxt(n);
     for (int h = sz(bv); h--;) {
+      int i = 0;
       vector<bool> b(n);
-      rep(i, 0, n) b[i] = (a[i] >> h) & 1;
+      ranges::stable_partition(a,
+        [&](ull x) { return b[i++] = (~x >> h) & 1; });
       bv[h] = b;
-      array it{begin(nxt), begin(nxt) + bv[h].cnt0(n)};
-      rep(i, 0, n) * it[b[i]]++ = a[i];
-      swap(a, nxt);
     }
   }
-  //! (k+1)th smallest number in [l,r)
-  //! kth(l,r,0) returns the min
-  //! @time O(log(max_val))
-  //! @space O(1)
   ull kth(int l, int r, int k) {
     ll res = 0;
     for (int h = sz(bv); h--;) {
-      int l0 = bv[h].cnt0(l), r0 = bv[h].cnt0(r);
+      int l0 = bv[h].cnt(l), r0 = bv[h].cnt(r);
       if (k < r0 - l0) l = l0, r = r0;
       else
         k -= r0 - l0, res |= 1ULL << h,
-          l += bv[h].cnt0(n) - l0, r += bv[h].cnt0(n) - r0;
+          l += bv[h].cnt(n) - l0, r += bv[h].cnt(n) - r0;
     }
     return res;
   }