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

📅 2026/7/10 4:42:12 👁️ 阅读次数 📝 编程学习
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` 取模,避免大数运算。