better stree impl

Author: NotFish232

src/rust_solutions/2179_count-good-triplets-in-an-array.rs

@@ -37,18 +37,23 @@ impl STree {
         }
     }
 
-    fn query(&mut self, idx_1: usize, idx_2: usize) -> i32 {
-        let (mut lp, mut rp) = (self.n + idx_1 - 1, self.n + idx_2);
+    fn query(&self, idx_1: usize, idx_2: usize) -> i32 {
+        let (mut lp, mut rp) = (self.n + idx_1 - 1, self.n + idx_2 - 1);
 
         let mut res = 0;
 
-        while lp < rp {
+       loop {
             if lp % 2 == 0 {
                 res += self.tree[lp];
                 lp += 1;
             }
-            if rp % 2 == 0 {
-                res += self.tree[rp - 1];
+            if rp % 2 == 1 {
+                res += self.tree[rp];
+                rp -= 1;
+            }
+
+            if lp > rp {
+                break;
             }
 
             lp = (lp - 1) / 2;
@@ -83,7 +88,7 @@ impl Solution {
             l_tree.update(pos, 1);
             r_tree.update(pos, 0);
 
-            let left = l_tree.query(0, pos - 1);
+            let left = if pos > 0 { l_tree.query(0, pos - 1) } else {0 };
             let right = r_tree.query(pos + 1, nums1.len() - 1);
             count += left as i64 * right as i64;
         }

src/rust_solutions/2411_smallest-subarrays-with-maximum-bitwise-or.rs

@@ -38,25 +38,30 @@ impl STree {
         }
     }
 
-    fn query(&self, l_idx: usize, r_idx: usize) -> i32 {
-        let (mut l_ptr, mut r_ptr) = (self.n + l_idx - 1, self.n + r_idx);
+    fn query(&self, idx_1: usize, idx_2: usize) -> i32 {
+        let (mut lp, mut rp) = (self.n + idx_1 - 1, self.n + idx_2 - 1);
 
-        let mut sum = 0;
+        let mut res = 0;
 
-        while l_ptr < r_ptr {
-            if l_ptr % 2 == 0 {
-                sum |= self.tree[l_ptr];
-                l_ptr += 1;
+        loop {
+            if lp % 2 == 0 {
+                res |= self.tree[lp];
+                lp += 1;
             }
-            if r_ptr % 2 == 0 {
-                sum |= self.tree[r_ptr - 1];
+            if rp % 2 == 1 {
+                res |= self.tree[rp];
+                rp -= 1;
             }
 
-            l_ptr = (l_ptr - 1) / 2;
-            r_ptr = (r_ptr - 1) / 2;
+            if lp > rp {
+                break;
+            }
+
+            lp = (lp - 1) / 2;
+            rp = (rp - 1) / 2;
         }
 
-        sum
+        res
     }
 }
 

src/rust_solutions/3355_zero-array-transformation-i.rs

@@ -40,20 +40,26 @@ impl STree {
     }
 
     fn query(&self, idx_1: usize, idx_2: usize) -> i32 {
-        let (mut l_ptr, mut r_ptr) = (self.n + idx_1 - 1, self.n + idx_2);
+        let (mut lp, mut rp) = (self.n + idx_1 - 1, self.n + idx_2 - 1);
+
         let mut res = 0;
 
-        while l_ptr < r_ptr {
-            if l_ptr % 2 == 0 {
-                res += self.v[l_ptr];
-                l_ptr += 1;
+       loop {
+            if lp % 2 == 0 {
+                res += self.v[lp];
+                lp += 1;
+            }
+            if rp % 2 == 1 {
+                res += self.v[rp];
+                rp -= 1;
             }
-            if r_ptr % 2 == 0 {
-                res += self.v[r_ptr - 1];
+
+            if lp > rp {
+                break;
             }
 
-            l_ptr = (l_ptr - 1) / 2;
-            r_ptr = (r_ptr - 1) / 2;
+            lp = (lp - 1) / 2;
+            rp = (rp - 1) / 2;
         }
 
         res