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

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

思路概述

- 对于每个查询，先单点修改 `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)`

完整代码

```python
class Node:
__slots__ = ('prod', 'cnt')
def __init__(self):
self.prod = 1 # 区间乘积 % k
self.cnt = [0] * 5 # cnt[m] = 非空前缀中乘积 % k == m 的个数

class SegTree:
def __init__(self, nums, k):
self.n = len(nums)
self.k = k
self.tree = [Node() for _ in range(self.n * 4)]
self.build(1, 0, self.n - 1, nums)

def build(self, node, l, r, nums):
if l == r:
v = nums[l] % self.k
self.tree[node].prod = v
self.tree[node].cnt = [0] * 5
self.tree[node].cnt[v] = 1
return
mid = (l + r) >> 1
self.build(node << 1, l, mid, nums)
self.build(node << 1 | 1, mid + 1, r, nums)
self.pull(node)

def merge(self, left, right):
res = Node()
res.prod = (left.prod * right.prod) % self.k

# 左区间的前缀直接保留
for i in range(self.k):
res.cnt[i] += left.cnt[i]

# 右区间的前缀需要乘上左区间的总积
for i in range(self.k):
if right.cnt[i]:
new_mod = (left.prod * i) % self.k
res.cnt[new_mod] += right.cnt[i]

return res

def pull(self, node):
self.tree[node] = self.merge(self.tree[node << 1], self.tree[node << 1 | 1])

def update(self, node, l, r, idx, val):
if l == r:
v = val % self.k
self.tree[node].prod = v
self.tree[node].cnt = [0] * 5
self.tree[node].cnt[v] = 1
return
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)

def query(self, node, l, r, ql, qr):
if ql <= l and r <= qr:
return self.tree[node]
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)
left = self.query(node << 1, l, mid, ql, qr)
right = self.query(node << 1 | 1, mid + 1, r, ql, qr)
return self.merge(left, right)

class Solution:
def resultArray(self, nums: List[int], k: int, queries: List[List[int]]) -> List[int]:
n = len(nums)
# 预处理：所有数先对 k 取模
nums = [x % k for x in nums]

seg = SegTree(nums, k)
ans = []

for q in queries:
idx, val, start, x = q
val %= k
seg.update(1, 0, n - 1, idx, val)
res = seg.query(1, 0, n - 1, start, n - 1)
ans.append(res.cnt[x])

return ans
```

关键点说明

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