Handling sum queries after update

Array Segment Tree

30.5% acceptance

Feb 25, 2026

198

You are given two 0-indexed arrays nums1 and nums2 and a 2D array queries of queries. There are three types of queries: For a query of type 1, queries[i] = [1, l, r]. Flip the values from 0 to 1 and from 1 to 0 in nums1 from index l to index r. Both l and r are 0-indexed. For a query of type 2, queries[i] = [2, p, 0]. For every index 0 <= i < n, set nums2[i] = nums2[i] + nums1[i] * p. For a query of type 3, queries[i] = [3, 0, 0]. Find the sum of the elements in nums2. Return an array containing all the answers to the third type queries.

Solution

Rust

Time O(2^n)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn handle_query(nums1: Vec<i32>, nums2: Vec<i32>, queries: Vec<Vec<i32>>) -> Vec<i64> {
    let n = nums1.len();
    let mut tree = vec![0i32; 4 * n];
    let mut lazy = vec![false; 4 * n];

    fn build(tree: &mut [i32], nums1: &[i32], node: usize, start: usize, end: usize) {
      if start == end {
        tree[node] = nums1[start];
        return;
      }
      let mid = (start + end) / 2;
      build(tree, nums1, 2 * node, start, mid);
      build(tree, nums1, 2 * node + 1, mid + 1, end);
      tree[node] = tree[2 * node] + tree[2 * node + 1];
    }

    fn push_down(tree: &mut [i32], lazy: &mut [bool], node: usize, start: usize, end: usize) {
      if lazy[node] {
        let mid = (start + end) / 2;
        let left = 2 * node;
        let right = 2 * node + 1;
        lazy[left] = !lazy[left];
        tree[left] = (mid - start + 1) as i32 - tree[left];
        lazy[right] = !lazy[right];
        tree[right] = (end - mid) as i32 - tree[right];
        lazy[node] = false;
      }
    }

    fn update(
      tree: &mut [i32], lazy: &mut [bool],
      node: usize, start: usize, end: usize, l: usize, r: usize,
    ) {
      if r < start || end < l { return; }
      if l <= start && end <= r {
        lazy[node] = !lazy[node];
        tree[node] = (end - start + 1) as i32 - tree[node];
        return;
      }
      push_down(tree, lazy, node, start, end);
      let mid = (start + end) / 2;
      update(tree, lazy, 2 * node, start, mid, l, r);
      update(tree, lazy, 2 * node + 1, mid + 1, end, l, r);
      tree[node] = tree[2 * node] + tree[2 * node + 1];
    }

    build(&mut tree, &nums1, 1, 0, n - 1);

    let mut sum2: i64 = nums2.iter().map(|&x| x as i64).sum();
    let mut result = Vec::new();

    for q in &queries {
      match q[0] {
        1 => update(&mut tree, &mut lazy, 1, 0, n - 1, q[1] as usize, q[2] as usize),
        2 => sum2 += tree[1] as i64 * q[1] as i64,
        3 => result.push(sum2),
        _ => {}
      }
    }
    result
  }
}