以下是 LeetCode 3525 求出数组的 X 值 II 的 Rust 实现,核心思路是线段树维护区间前缀积的模分布。

思路概述

- 对于每个查询,先单点修改 `nums[index]`,再移除前缀保留 `[start..]`,最后可以删除任意后缀(即保留任意非空前缀)。
- 因此每个查询的答案等价于:在子数组 `nums[start..]` 的所有非空前缀中,乘积模 `k` 等于 `x` 的个数。
- 由于 `k ≤ 5`,线段树每个节点只需维护:
  - `prod`:区间整体乘积模 `k`
  - `cnt[5]`:`cnt[m]` 表示该区间内,有多少个非空前缀的乘积模 `k` 等于 `m`
- 合并两个子区间时,右子树的前缀需要乘上左子树的总积再模 `k`。

复杂度

- 时间:`O((n + q) · k² · log n)`,由于 `k ≤ 5`,实际约为 `O((n + q) log n)`
- 空间:`O(n)`

完整代码

```rust
#[derive(Clone, Copy)]
struct Node {
    prod: i32,      // 区间乘积 % k
    cnt: [i32; 5],  // cnt[m] = 非空前缀中乘积 % k == m 的个数
}

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

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

    fn build(&mut self, node: usize, l: usize, r: usize, nums: &[i32]) {
        if l == r {
            let v = nums[l] % self.k;
            self.tree[node].prod = v;
            self.tree[node].cnt = [0; 5];
            self.tree[node].cnt[v as usize] = 1;
            return;
        }
        let mid = (l + r) >> 1;
        self.build(node << 1, l, mid, nums);
        self.build(node << 1 | 1, mid + 1, r, nums);
        self.pull(node);
    }

    // 合并左右两个区间
    fn merge(&self, left: &Node, right: &Node) -> Node {
        let mut cnt = [0; 5];
        let k = self.k as usize;

        // 左区间的前缀直接保留
        for i in 0..k {
            cnt[i] += left.cnt[i];
        }

        // 右区间的前缀需要乘上左区间的总积
        for i in 0..k {
            if right.cnt[i] != 0 {
                let new_mod = (left.prod * i as i32) % self.k;
                cnt[new_mod as usize] += right.cnt[i];
            }
        }

        Node {
            prod: (left.prod * right.prod) % self.k,
            cnt,
        }
    }

    fn pull(&mut self, node: usize) {
        let left = self.tree[node << 1];
        let right = self.tree[node << 1 | 1];
        self.tree[node] = self.merge(&left, &right);
    }

    fn update(&mut self, node: usize, l: usize, r: usize, idx: usize, val: i32) {
        if l == r {
            let v = val % self.k;
            self.tree[node].prod = v;
            self.tree[node].cnt = [0; 5];
            self.tree[node].cnt[v as usize] = 1;
            return;
        }
        let mid = (l + r) >> 1;
        if idx <= mid {
            self.update(node << 1, l, mid, idx, val);
        } else {
            self.update(node << 1 | 1, mid + 1, r, idx, val);
        }
        self.pull(node);
    }

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

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

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

            seg.update(1, 0, n - 1, idx, val);
            let res = seg.query(1, 0, n - 1, start, n - 1);
            ans.push(res.cnt[x]);
        }

        ans
    }
}
```

关键点说明

1. `cnt` 的含义:每个节点维护的是从该区间左端点开始的所有非空前缀的模分布。这样查询任意子区间时,递归返回的子节点天然满足“从左端点开始”的要求,合并时只需把右子树的前缀乘上左子树的总积即可。
2. `k ≤ 5` 的利用:`cnt` 数组固定开 5,实际只用前 `k` 个;合并时的双重循环最多 25 次运算,常数极小。
3. 单点更新:修改叶子后自底向上 `pull`,保持每个节点的 `prod` 和 `cnt` 正确。

 

Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐