wavelet nit (#172)

* Optimize wavelet_matrix constructor logic

* [auto-verifier] verify commit fcb8967

* Update wavelet_matrix.hpp

* [auto-verifier] verify commit cd5257c

* Revert "Update wavelet_matrix.hpp"

This reverts commit cd5257c.

* fix

* fix

* fix

* another fix

* [auto-verifier] verify commit 45dd8c6

* another nit

---------

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

Author: lrvideckis

.verify-helper/timestamps.remote.json

@@ -25,7 +25,7 @@
 "tests/library_checker_aizu_tests/data_structures/implicit_seg_tree.test.cpp": "2026-01-18 11:15:41 +0000",
 "tests/library_checker_aizu_tests/data_structures/kruskal_tree_aizu.test.cpp": "2026-01-18 02:20:40 +0000",
 "tests/library_checker_aizu_tests/data_structures/kth_smallest_pst.test.cpp": "2026-01-18 11:15:41 +0000",
-"tests/library_checker_aizu_tests/data_structures/kth_smallest_wavelet_matrix.test.cpp": "2026-01-18 02:20:40 +0000",
+"tests/library_checker_aizu_tests/data_structures/kth_smallest_wavelet_matrix.test.cpp": "2026-01-19 02:49:20 -0700",
 "tests/library_checker_aizu_tests/data_structures/lazy_segment_tree.test.cpp": "2026-01-18 04:18:37 -0700",
 "tests/library_checker_aizu_tests/data_structures/lazy_segment_tree_constructor.test.cpp": "2026-01-18 04:18:37 -0700",
 "tests/library_checker_aizu_tests/data_structures/lazy_segment_tree_inc.test.cpp": "2026-01-18 11:04:58 +0000",
@@ -75,7 +75,7 @@
 "tests/library_checker_aizu_tests/handmade_tests/functional_graph.test.cpp": "2025-08-06 16:18:37 -0600",
 "tests/library_checker_aizu_tests/handmade_tests/hilbert_mos.test.cpp": "2026-01-18 02:20:40 +0000",
 "tests/library_checker_aizu_tests/handmade_tests/manacher.test.cpp": "2026-01-18 11:15:41 +0000",
-"tests/library_checker_aizu_tests/handmade_tests/merge_st_and_wavelet.test.cpp": "2026-01-18 11:04:58 +0000",
+"tests/library_checker_aizu_tests/handmade_tests/merge_st_and_wavelet.test.cpp": "2026-01-19 02:51:03 -0700",
 "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_division.test.cpp": "2025-09-07 16:12:35 -0600",
 "tests/library_checker_aizu_tests/handmade_tests/n_choose_k.test.cpp": "2025-08-28 13:19:16 -0600",

library/data_structures_[l,r)/seg_tree_uncommon/wavelet_count_less.hpp

@@ -1,8 +1,8 @@
 #pragma once
 //! count of i in [l..r) such that a[i] < ub
-//! @time O(log(max_val))
+//! @time O(lg)
 //! @space O(1)
-int count(int l, int r, ull ub) {
+int count(int l, int r, ll ub) {
   int res = 0;
   for (int h = sz(bv); h--;) {
     int l0 = bv[h].cnt(l), r0 = bv[h].cnt(r);

library/data_structures_[l,r)/seg_tree_uncommon/wavelet_matrix.hpp

@@ -1,27 +1,27 @@
 #pragma once
 #include "wavelet_bit_vec.hpp"
 //! @code
-//!   vector<ull> a(n);
-//!   wavelet_matrix wm(a, 1e9); // requires a[i] <= 1e9
+//!   vector<ll> a(n);
+//!   wavelet_matrix wm(a, 30); // 0 <= a[i] < (1LL<<30)
 //!   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)
+//! @time O(n * lg + q * lg)
+//! @space O(n * lg / 64)
 struct wavelet_matrix {
   int n;
   vector<bit_vec> bv;
-  wavelet_matrix(vector<ull> a, ull max_val):
-    n(sz(a)), bv(bit_width(max_val), {{}}) {
+  wavelet_matrix(vector<ll> a, int lg):
+    n(sz(a)), bv(lg, {{}}) {
     for (int h = sz(bv); h--;) {
-      int i = 0;
       vector<bool> b(n);
-      ranges::stable_partition(a,
-        [&](ull x) { return b[i++] = (~x >> h) & 1; });
+      rep(i, 0, n) b[i] = (~a[i] >> h) & 1;
       bv[h] = b;
+      ranges::stable_partition(a,
+        [&](ll x) { return (~x >> h) & 1; });
     }
   }
-  ull kth(int l, int r, int k) {
+  ll kth(int l, int r, int k) {
     ll res = 0;
     for (int h = sz(bv); h--;) {
       int l0 = bv[h].cnt(l), r0 = bv[h].cnt(r);

tests/library_checker_aizu_tests/data_structures/kth_smallest_wavelet_matrix.test.cpp

@@ -6,12 +6,12 @@ int main() {
   cin.tie(0)->sync_with_stdio(0);
   int n, q;
   cin >> n >> q;
-  vector<ull> a(n);
+  vector<ll> a(n);
   for (int i = 0; i < n; i++) {
     cin >> a[i];
     a[i] <<= 23;
   }
-  wavelet_matrix wm(a, 1'000'000'000LL << 23);
+  wavelet_matrix wm(a, 30 + 23);
   while (q--) {
     int l, r, k;
     cin >> l >> r >> k;

tests/library_checker_aizu_tests/handmade_tests/merge_st_and_wavelet.test.cpp

@@ -16,9 +16,9 @@ int main() {
       generate(begin(arr), end(arr),
         [&]() { return rnd<int>(minn, maxn); });
       merge_sort_tree mst(arr);
-      vector<ull> arr_shifted(n);
+      vector<ll> arr_shifted(n);
       rep(i, 0, n) arr_shifted[i] = arr[i] - minn;
-      wavelet_matrix wm(arr_shifted, maxn - minn);
+      wavelet_matrix wm(arr_shifted, 11);
       for (int queries = 30; queries--;) {
         int x = rnd<int>(minn, maxn);
         int y = rnd<int>(minn, maxn);