Kimi LeetCode 3525. 求出数组的 X 值 II Go实现

以下是 LeetCode 3525 求出数组的 X 值 II 的 Go 实现核心思路是线段树维护区间前缀积的模分布。思路概述- 对于每个查询先单点修改 nums[index]再移除前缀保留 [start..]最后可以删除任意后缀即保留任意非空前缀。- 因此每个查询的答案等价于在子数组 nums[start..] 的所有非空前缀中乘积模 k 等于 x 的个数。- 由于 k ≤ 5线段树每个节点只需维护- prod区间整体乘积模 k- cntcnt[m] 表示该区间内有多少个非空前缀的乘积模 k 等于 m- 合并两个子区间时右子树的前缀需要乘上左子树的总积再模 k。复杂度- 时间O((n q) · k² · log n)由于 k ≤ 5实际约为 O((n q) log n)- 空间O(n)完整代码gotype Node struct {prod int // 区间乘积 % kcnt [5]int // cnt[m] 非空前缀中乘积 % k m 的个数}type SegTree struct {k inttree []Noden int}func NewSegTree(nums []int, k int) *SegTree {n : len(nums)seg : SegTree{k: k,tree: make([]Node, n*4),n: n,}seg.build(1, 0, n-1, nums)return seg}func (s *SegTree) build(node, l, r int, nums []int) {if l r {v : nums[l] % s.ks.tree[node].prod vs.tree[node].cnt [5]int{}s.tree[node].cnt[v] 1return}mid : (l r) 1s.build(node1, l, mid, nums)s.build(node1|1, mid1, r, nums)s.pull(node)}func (s *SegTree) merge(left, right Node) Node {var res Noderes.prod (left.prod * right.prod) % s.k// 左区间的前缀直接保留for i : 0; i s.k; i {res.cnt[i] left.cnt[i]}// 右区间的前缀需要乘上左区间的总积for i : 0; i s.k; i {if right.cnt[i] ! 0 {newMod : (left.prod * i) % s.kres.cnt[newMod] right.cnt[i]}}return res}func (s *SegTree) pull(node int) {s.tree[node] s.merge(s.tree[node1], s.tree[node1|1])}func (s *SegTree) update(node, l, r, idx, val int) {if l r {v : val % s.ks.tree[node].prod vs.tree[node].cnt [5]int{}s.tree[node].cnt[v] 1return}mid : (l r) 1if idx mid {s.update(node1, l, mid, idx, val)} else {s.update(node1|1, mid1, r, idx, val)}s.pull(node)}func (s *SegTree) query(node, l, r, ql, qr int) Node {if ql l r qr {return s.tree[node]}mid : (l r) 1if qr mid {return s.query(node1, l, mid, ql, qr)}if ql mid {return s.query(node1|1, mid1, r, ql, qr)}left : s.query(node1, l, mid, ql, qr)right : s.query(node1|1, mid1, r, ql, qr)return s.merge(left, right)}func resultArray(nums []int, k int, queries [][]int) []int {n : len(nums)for i : range nums {nums[i] % k}seg : NewSegTree(nums, k)ans : make([]int, 0, len(queries))for _, q : range queries {idx, val, start, x : q[0], q[1]%k, q[2], q[3]seg.update(1, 0, n-1, idx, val)res : seg.query(1, 0, n-1, start, n-1)ans append(ans, res.cnt[x])}return ans}关键点说明1. cnt 的含义每个节点维护的是从该区间左端点开始的所有非空前缀的模分布。这样查询任意子区间时递归返回的子节点天然满足从左端点开始的要求合并时只需把右子树的前缀乘上左子树的总积即可。2. k ≤ 5 的利用cnt 数组固定开 5实际只用前 k 个合并时的双重循环最多 25 次运算常数极小。3. 单点更新修改叶子后自底向上 pull保持每个节点的 prod 和 cnt 正确。4. 预处理取模建树前和更新时都把数值对 k 取模避免大数运算。