Alternating groups iii

impl Solution {
  pub fn number_of_alternating_groups(colors: Vec<i32>, queries: Vec<Vec<i32>>) -> Vec<i32> {
    let n = colors.len();
    let mut cols = colors;

    // ok[i] = 1 iff cols[i] != cols[(i+1)%n] (edge between tile i and i+1 is "alternating")
    let mut ok: Vec<i32> = (0..n).map(|i| (cols[i] != cols[(i + 1) % n]) as i32).collect();

    // zeros = set of positions where ok[i] == 0
    let mut zeros: std::collections::BTreeSet<usize> =
      (0..n).filter(|&i| ok[i] == 0).collect();

    // rlen(z1, z2, n): number of positions strictly between z1 and z2 going forward
    fn rlen(z1: usize, z2: usize, n: usize) -> usize {
      (z2 + 2 * n - z1 - 1) % n
    }

    // Fenwick trees indexed by run length (1..=n)
    let mut bit_c = vec![0i64; n + 2]; // count of runs by length
    let mut bit_s = vec![0i64; n + 2]; // sum of lengths

    fn bit_upd(bit: &mut Vec<i64>, mut i: usize, d: i64) {
      i += 1;
      while i < bit.len() {
        bit[i] += d;
        i += i & i.wrapping_neg();
      }
    }
    fn bit_qry(bit: &Vec<i64>, mut i: usize) -> i64 {
      i += 1;
      let mut s = 0i64;
      while i > 0 {
        s += bit[i];
        i -= i & i.wrapping_neg();
      }
      s
    }
    fn add_run(bc: &mut Vec<i64>, bs: &mut Vec<i64>, l: usize, d: i64) {
      if l == 0 {
        return;
      }
      bit_upd(bc, l, d);
      bit_upd(bs, l, d * l as i64);
    }

    // Initialize BIT from current zeros
    if zeros.is_empty() {
      add_run(&mut bit_c, &mut bit_s, n, 1);
    } else {
      let zv: Vec<usize> = zeros.iter().copied().collect();
      let m = zv.len();
      for i in 0..m {
        let l = rlen(zv[i], zv[(i + 1) % m], n);
        add_run(&mut bit_c, &mut bit_s, l, 1);
      }
    }

    let mut result = Vec::new();

    for q in &queries {
      if q[0] == 1 {
        let k = q[1] as usize;
        let ans = if zeros.is_empty() {
          n as i32
        } else {
          let min_r = k - 1; // run must have length >= min_r
          if min_r > n {
            0
          } else {
            let tot_c = bit_qry(&bit_c, n);
            let tot_s = bit_qry(&bit_s, n);
            let pr_c = if min_r > 1 { bit_qry(&bit_c, min_r - 1) } else { 0 };
            let pr_s = if min_r > 1 { bit_qry(&bit_s, min_r - 1) } else { 0 };
            let su_c = tot_c - pr_c;
            let su_s = tot_s - pr_s;
            // contribution = sum_{L>=min_r} (L - min_r + 1) = su_s - (min_r-1)*su_c
            (su_s - (min_r as i64 - 1) * su_c) as i32
          }
        };
        result.push(ans);
      } else {
        let idx = q[1] as usize;
        let new_color = q[2];
        if cols[idx] == new_color {
          continue;
        }
        cols[idx] = new_color;

        // Update ok edge to the left: ok[(idx-1+n)%n] = (cols[idx-1] != cols[idx])
        let left = (idx + n - 1) % n;
        let new_ok_l = (cols[left] != cols[idx]) as i32;
        if ok[left] != new_ok_l {
          Self::toggle_ok(
            left, new_ok_l, &mut ok, &mut zeros, &mut bit_c, &mut bit_s, n,
          );
        }

        // Update ok edge at idx: ok[idx] = (cols[idx] != cols[(idx+1)%n])
        let right = (idx + 1) % n;
        let new_ok_r = (cols[idx] != cols[right]) as i32;
        if ok[idx] != new_ok_r {
          Self::toggle_ok(
            idx, new_ok_r, &mut ok, &mut zeros, &mut bit_c, &mut bit_s, n,
          );
        }
      }
    }

    result
  }

  fn toggle_ok(
    j: usize,
    new_val: i32,
    ok: &mut Vec<i32>,
    zeros: &mut std::collections::BTreeSet<usize>,
    bc: &mut Vec<i64>,
    bs: &mut Vec<i64>,
    n: usize,
  ) {
    fn rlen(z1: usize, z2: usize, n: usize) -> usize {
      (z2 + 2 * n - z1 - 1) % n
    }
    fn bit_upd(bit: &mut Vec<i64>, mut i: usize, d: i64) {
      i += 1;
      while i < bit.len() {
        bit[i] += d;
        i += i & i.wrapping_neg();
      }
    }
    fn add_run(bc: &mut Vec<i64>, bs: &mut Vec<i64>, l: usize, d: i64) {
      if l == 0 {
        return;
      }
      bit_upd(bc, l, d);
      bit_upd(bs, l, d * l as i64);
    }

    if new_val == 0 {
      // 1 → 0: split the run containing j
      if zeros.is_empty() {
        // single circular run of n → split into run of n-1
        add_run(bc, bs, n, -1);
        add_run(bc, bs, n - 1, 1);
      } else {
        let prev_z = zeros
          .range(..j)
          .next_back()
          .copied()
          .unwrap_or_else(|| *zeros.iter().next_back().unwrap());
        let next_z = zeros
          .range((j + 1)..)
          .next()
          .copied()
          .unwrap_or_else(|| *zeros.iter().next().unwrap());
        let old_len = rlen(prev_z, next_z, n);
        let left_len = rlen(prev_z, j, n);
        let right_len = rlen(j, next_z, n);
        add_run(bc, bs, old_len, -1);
        add_run(bc, bs, left_len, 1);
        add_run(bc, bs, right_len, 1);
      }
      zeros.insert(j);
    } else {
      // 0 → 1: merge runs on both sides of j
      if zeros.len() == 1 {
        // Only zero is j; removing it makes full circle
        add_run(bc, bs, n - 1, -1); // rlen(j, j) = n-1
        zeros.remove(&j);
        add_run(bc, bs, n, 1);
      } else {
        let prev_z = zeros
          .range(..j)
          .next_back()
          .copied()
          .unwrap_or_else(|| *zeros.iter().next_back().unwrap());
        let next_z = zeros
          .range((j + 1)..)
          .next()
          .copied()
          .unwrap_or_else(|| *zeros.iter().next().unwrap());
        let left_len = rlen(prev_z, j, n);
        let right_len = rlen(j, next_z, n);
        let new_len = rlen(prev_z, next_z, n);
        add_run(bc, bs, left_len, -1);
        add_run(bc, bs, right_len, -1);
        zeros.remove(&j);
        add_run(bc, bs, new_len, 1);
      }
    }
    ok[j] = new_val;
  }
}