Kimi LeetCode 3525. 求出数组的 X 值 II Go实现
以下是 LeetCode 3525 求出数组的 X 值 II 的 Go 实现，核心思路是线段树维护区间前缀积的模分布。
思路概述
- 对于每个查询，先单点修改 `nums[index]`，再移除前缀保留 `[start..]`，最后可以删除任意后缀（即保留任意非空前缀）。
- 因此每个查询的答案等价于：在子数组 `nums[start..]` 的所有非空前缀中，乘积模 `k` 等于 `x` 的个数。
- 由于 `k ≤ 5`，线段树每个节点只需维护：
- `prod`：区间整体乘积模 `k`
- `cnt`：`cnt[m]` 表示该区间内，有多少个非空前缀的乘积模 `k` 等于 `m`
- 合并两个子区间时，右子树的前缀需要乘上左子树的总积再模 `k`。
复杂度
- 时间：`O((n + q) · k² · log n)`，由于 `k ≤ 5`，实际约为 `O((n + q) log n)`
- 空间：`O(n)`
完整代码
```go
type Node struct {
prod int // 区间乘积 % k
cnt [5]int // cnt[m] = 非空前缀中乘积 % k == m 的个数
}
type SegTree struct {
k int
tree []Node
n 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.k
s.tree[node].prod = v
s.tree[node].cnt = [5]int{}
s.tree[node].cnt[v] = 1
return
}
mid := (l + r) >> 1
s.build(node<<1, l, mid, nums)
s.build(node<<1|1, mid+1, r, nums)
s.pull(node)
}
func (s *SegTree) merge(left, right Node) Node {
var res Node
res.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.k
res.cnt[newMod] += right.cnt[i]
}
}
return res
}
func (s *SegTree) pull(node int) {
s.tree[node] = s.merge(s.tree[node<<1], s.tree[node<<1|1])
}
func (s *SegTree) update(node, l, r, idx, val int) {
if l == r {
v := val % s.k
s.tree[node].prod = v
s.tree[node].cnt = [5]int{}
s.tree[node].cnt[v] = 1
return
}
mid := (l + r) >> 1
if idx <= mid {
s.update(node<<1, l, mid, idx, val)
} else {
s.update(node<<1|1, mid+1, 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) >> 1
if qr <= mid {
return s.query(node<<1, l, mid, ql, qr)
}
if ql > mid {
return s.query(node<<1|1, mid+1, r, ql, qr)
}
left := s.query(node<<1, l, mid, ql, qr)
right := s.query(node<<1|1, mid+1, 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` 取模，避免大数运算。