Find x value of array ii

Array Math Segment Tree

30.6% acceptance

Feb 25, 2026

You are given an array of positive integers nums and a positive integer k. You are also given a 2D array queries, where queries[i] = [indexi, valuei, starti, xi]. Update nums[indexi] to valuei (persists), remove prefix up to starti-1, count ways to remove a suffix so the product of remaining elements ≡ xi (mod k). Return result[i] for each query.

Solution

Rust

Time O(n²)

Space O(n)

LeetCode

solution.rs

// Segment tree node:
// cnt[in][out] = number of positions in [l,r] where product(nums[l..=end]) starting
//                with incoming product `in` equals `out` (mod k)
// full[in]     = (in * product(nums[l..=r])) % k  (product of entire range)
#[derive(Clone)]
struct Node {
  cnt: [[i32; 5]; 5],
  full: [usize; 5],
}

impl Node {
  fn identity() -> Self {
    Node { cnt: [[0; 5]; 5], full: [0, 1, 2, 3, 4] }
  }
  fn leaf(v: usize, k: usize) -> Self {
    let mut node = Node { cnt: [[0; 5]; 5], full: [0; 5] };
    for inp in 0..k {
      let out = (inp * v) % k;
      node.cnt[inp][out] += 1;
      node.full[inp] = out;
    }
    node
  }
  fn merge(l: &Node, r: &Node, k: usize) -> Self {
    let mut node = Node { cnt: [[0; 5]; 5], full: [0; 5] };
    for inp in 0..k {
      // endings in left child
      for out in 0..k {
        node.cnt[inp][out] += l.cnt[inp][out];
      }
      // endings in right child: incoming product is l.full[inp]
      let mid = l.full[inp];
      for out in 0..k {
        node.cnt[inp][out] += r.cnt[mid][out];
      }
      node.full[inp] = r.full[l.full[inp]];
    }
    node
  }
}

struct SegTree {
  n: usize,
  k: usize,
  tree: Vec<Node>,
}

impl SegTree {
  fn build(nums: &[i32], k: usize) -> Self {
    let n = nums.len();
    let mut tree = vec![Node::identity(); 4 * n];
    Self::build_rec(&mut tree, nums, k, 1, 0, n - 1);
    SegTree { n, k, tree }
  }

  fn build_rec(tree: &mut Vec<Node>, nums: &[i32], k: usize, node: usize, l: usize, r: usize) {
    if l == r {
      tree[node] = Node::leaf(nums[l] as usize % k, k);
      return;
    }
    let mid = (l + r) / 2;
    Self::build_rec(tree, nums, k, 2 * node, l, mid);
    Self::build_rec(tree, nums, k, 2 * node + 1, mid + 1, r);
    let (left, right) = (tree[2 * node].clone(), tree[2 * node + 1].clone());
    tree[node] = Node::merge(&left, &right, k);
  }

  fn update(&mut self, pos: usize, val: usize) {
    self.update_rec(1, 0, self.n - 1, pos, val);
  }

  fn update_rec(&mut self, node: usize, l: usize, r: usize, pos: usize, val: usize) {
    if l == r {
      self.tree[node] = Node::leaf(val % self.k, self.k);
      return;
    }
    let mid = (l + r) / 2;
    if pos <= mid {
      self.update_rec(2 * node, l, mid, pos, val);
    } else {
      self.update_rec(2 * node + 1, mid + 1, r, pos, val);
    }
    let (left, right) = (self.tree[2 * node].clone(), self.tree[2 * node + 1].clone());
    self.tree[node] = Node::merge(&left, &right, self.k);
  }

  // Query range [ql, qr], return cnt[1 % k][x]
  fn query(&self, ql: usize, qr: usize, x: usize) -> i32 {
    self.query_rec(1, 0, self.n - 1, ql, qr).cnt[1 % self.k][x]
  }

  fn query_rec(&self, node: usize, l: usize, r: usize, ql: usize, qr: usize) -> Node {
    if ql <= l && r <= qr {
      return self.tree[node].clone();
    }
    let mid = (l + r) / 2;
    if qr <= mid {
      return self.query_rec(2 * node, l, mid, ql, qr);
    }
    if ql > mid {
      return self.query_rec(2 * node + 1, mid + 1, r, ql, qr);
    }
    let left = self.query_rec(2 * node, l, mid, ql, qr);
    let right = self.query_rec(2 * node + 1, mid + 1, r, ql, qr);
    Node::merge(&left, &right, self.k)
  }
}

impl Solution {
  pub fn result_array(mut nums: Vec<i32>, k: i32, queries: Vec<Vec<i32>>) -> Vec<i32> {
    let k = k as usize;
    let n = nums.len();
    let mut seg = SegTree::build(&nums, k);
    let mut result = Vec::with_capacity(queries.len());

    for q in &queries {
      let idx = q[0] as usize;
      let val = q[1] as usize;
      let start = q[2] as usize;
      let x = q[3] as usize;

      nums[idx] = val as i32;
      seg.update(idx, val);

      result.push(seg.query(start, n - 1, x));
    }

    result
  }
}