Maximize count of distinct primes after split

use std::collections::BTreeSet;

struct SegTree {
  n: usize,
  tree: Vec<i32>,
  lazy: Vec<i32>,
}

impl SegTree {
  fn new(n: usize) -> Self {
    let size = 4 * (n + 1);
    SegTree { n, tree: vec![0; size], lazy: vec![0; size] }
  }

  fn push_down(&mut self, node: usize) {
    let v = self.lazy[node];
    if v != 0 {
      for &c in &[2 * node, 2 * node + 1] {
        self.tree[c] += v;
        self.lazy[c] += v;
      }
      self.lazy[node] = 0;
    }
  }

  fn range_add(&mut self, l: usize, r: usize, val: i32) {
    if l > r { return; }
    let n = self.n;
    self.range_add_inner(1, 0, n - 1, l, r, val);
  }

  fn range_add_inner(&mut self, node: usize, nl: usize, nr: usize, l: usize, r: usize, val: i32) {
    if l > nr || r < nl { return; }
    if l <= nl && nr <= r {
      self.tree[node] += val;
      self.lazy[node] += val;
      return;
    }
    self.push_down(node);
    let mid = (nl + nr) / 2;
    self.range_add_inner(2 * node, nl, mid, l, r, val);
    self.range_add_inner(2 * node + 1, mid + 1, nr, l, r, val);
    self.tree[node] = self.tree[2 * node].max(self.tree[2 * node + 1]);
  }

  fn max_query(&mut self, l: usize, r: usize) -> i32 {
    if l > r { return 0; }
    let n = self.n;
    self.max_query_inner(1, 0, n - 1, l, r)
  }

  fn max_query_inner(&mut self, node: usize, nl: usize, nr: usize, l: usize, r: usize) -> i32 {
    if l > nr || r < nl { return 0; }
    if l <= nl && nr <= r { return self.tree[node]; }
    self.push_down(node);
    let mid = (nl + nr) / 2;
    self.max_query_inner(2 * node, nl, mid, l, r)
      .max(self.max_query_inner(2 * node + 1, mid + 1, nr, l, r))
  }
}

impl Solution {
  pub fn maximum_count(mut nums: Vec<i32>, queries: Vec<Vec<i32>>) -> Vec<i32> {
    const LIMIT: usize = 100001;
    let mut is_prime = vec![true; LIMIT];
    is_prime[0] = false;
    is_prime[1] = false;
    let mut i = 2;
    while i * i < LIMIT {
      if is_prime[i] {
        let mut j = i * i;
        while j < LIMIT {
          is_prime[j] = false;
          j += i;
        }
      }
      i += 1;
    }

    let n = nums.len();
    // For each prime value, maintain a sorted set of all positions it occupies.
    let mut positions: Vec<BTreeSet<usize>> = vec![BTreeSet::new(); LIMIT];
    let mut seg = SegTree::new(n);
    let mut distinct_primes = 0i32;

    // Activate the current (min,max) range of prime v in the seg tree.
    let activate = |positions: &Vec<BTreeSet<usize>>, seg: &mut SegTree, v: usize, delta: i32| {
      let s = &positions[v];
      if s.len() >= 2 {
        let mn = *s.iter().next().unwrap();
        let mx = *s.iter().next_back().unwrap();
        seg.range_add(mn + 1, mx, delta);
      }
    };

    // Build initial state
    for (idx, &v) in nums.iter().enumerate() {
      let v = v as usize;
      if v < LIMIT && is_prime[v] {
        activate(&positions, &mut seg, v, -1); // remove old range (harmless if empty)
        let was_empty = positions[v].is_empty();
        positions[v].insert(idx);
        if was_empty { distinct_primes += 1; }
        activate(&positions, &mut seg, v, 1);
      }
    }

    let mut result = Vec::with_capacity(queries.len());

    for q in &queries {
      let idx = q[0] as usize;
      let new_val = q[1] as usize;
      let old_val = nums[idx] as usize;

      // Remove old value at idx
      if old_val < LIMIT && is_prime[old_val] {
        activate(&positions, &mut seg, old_val, -1);
        positions[old_val].remove(&idx);
        if positions[old_val].is_empty() {
          distinct_primes -= 1;
        } else {
          activate(&positions, &mut seg, old_val, 1);
        }
      }

      nums[idx] = new_val as i32;

      // Add new value at idx
      if new_val < LIMIT && is_prime[new_val] {
        activate(&positions, &mut seg, new_val, -1);
        let was_empty = positions[new_val].is_empty();
        positions[new_val].insert(idx);
        if was_empty { distinct_primes += 1; }
        activate(&positions, &mut seg, new_val, 1);
      }

      let bonus = seg.max_query(1, n - 1);
      result.push(distinct_primes + bonus);
    }

    result
  }
}