Bit walk (#88)

* bit walk

* add example

* different interface actually

* fix docs

* Update fenwick.rs

---------

Co-authored-by: Luke Videckis <lukevideckis@gmail.com>

Author: lrvideckis

Cargo.toml

@@ -216,6 +216,10 @@ path = "examples/monotonic/count_rects.rs"
 name = "cartesian_tree"
 path = "examples/monotonic/cartesian_tree.rs"
 
+[[example]]
+name = "fenwick_kth"
+path = "examples/data_structures/fenwick_kth.rs"
+
 [[example]]
 name = "disjoint_rmq_non_commutative"
 path = "examples/data_structures/disjoint_rmq_non_commutative.rs"

examples/data_structures/fenwick_kth.rs

@@ -0,0 +1,54 @@
+// verification-helper: PROBLEM https://judge.yosupo.jp/problem/predecessor_problem
+
+use proconio::input;
+use programming_team_code_rust::data_structures::fenwick::Fenwick;
+
+fn main() {
+    input! {
+        n: usize,
+        q: usize,
+        s: String,
+    }
+
+    let mut fenwick =
+        Fenwick::<i32>::build_on_array(&s.chars().map(|c| (c == '1') as i32).collect::<Vec<_>>());
+
+    for _ in 0..q {
+        input! {
+            t: u8,
+            k: usize,
+        }
+        match t {
+            0 => {
+                if fenwick.sum(k..k + 1) == 0 {
+                    fenwick.add(k, 1);
+                }
+            }
+            1 => {
+                if fenwick.sum(k..k + 1) == 1 {
+                    fenwick.add(k, -1);
+                }
+            }
+            2 => {
+                println!("{}", fenwick.sum(k..k + 1));
+            }
+            3 => {
+                let cnt = fenwick.sum(0..k);
+                let res = fenwick.kth(cnt + 1);
+                if res == n {
+                    println!("-1");
+                } else {
+                    println!("{}", res);
+                }
+            }
+            _ => {
+                let cnt_le = fenwick.sum(0..k + 1);
+                if cnt_le == 0 {
+                    println!("-1");
+                } else {
+                    println!("{}", fenwick.kth(cnt_le));
+                }
+            }
+        }
+    }
+}

src/data_structures/fenwick.rs

@@ -9,6 +9,8 @@
 /// fenwick.add(2, 2);
 /// fenwick.add(3, 1);
 /// assert_eq!(fenwick.sum(1..3), 5);
+/// assert_eq!(fenwick.kth(5), 2);
+/// assert!(std::panic::catch_unwind(|| fenwick.kth(0)).is_err());
 /// ```
 pub struct Fenwick<T> {
     ary: Vec<T>,
@@ -75,4 +77,28 @@ impl<T: Clone + Default + std::ops::AddAssign<T>> Fenwick<T> {
     {
         self.accum(range.end) - self.accum(range.start)
     }
+
+    /// Gets maximum pos such that sum of [0, pos) < sum
+    ///
+    /// Requires fenwick.sum(i..i + 1) >= 0 and sum > 0
+    ///
+    /// # Complexity
+    /// - Time: O(log n)
+    /// - Space: O(1)
+    pub fn kth(&self, mut sum: T) -> usize
+    where
+        T: std::ops::SubAssign<T> + std::cmp::PartialOrd,
+    {
+        assert!(sum > T::default());
+        let mut pos = 0;
+        let mut pw = self.ary.len().next_power_of_two();
+        while pw > 0 {
+            if pos + pw <= self.ary.len() && self.ary[pos + pw - 1] < sum {
+                pos += pw;
+                sum -= self.ary[pos - 1].clone();
+            }
+            pw /= 2;
+        }
+        pos
+    }
 }