Kimi LeetCode 3901. 好子序列查询 Rust实现

根据搜索结果，LeetCode 3901「好子序列查询」的核心思路是：线段树 + GCD。下面给出 Rust 实现。

---

思路

1. 只关心 p 的倍数：如果 `nums[i]` 不是 `p` 的倍数，它对 GCD 为 `p` 的子序列没有贡献，在线段树中视为 `0`。
2. 线段树维护所有 p 倍数的 GCD：记为 `g`。
- 若 `g != p`，则不存在好子序列。
- 若 `g == p`，则所有 p 的倍数整体 GCD 已是 `p`。
3. 长度严格小于 n：
- 若 `cnt < n`（并非所有元素都是 p 的倍数），直接取所有 p 的倍数即可，长度必然 `< n`。
- 若 `cnt == n`（所有元素都是 p 的倍数），需要删一个元素后 GCD 仍为 `p`。
- 关键结论：若 `n > 6` 且整体 GCD 为 `p`，则一定可以删一个元素使 GCD 仍为 `p`。因此只需在 `n <= 6` 时暴力枚举删除位置。
4. 线段树操作：单点更新、区间查询 GCD。

---

Rust 实现

```rust
use std::cmp::min;

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

impl SegTree {
fn new(n: usize) -> Self {
Self {
n,
tree: vec![0; n * 4],
}
}

fn build(&mut self, u: usize, l: usize, r: usize, arr: &[i32]) {
if l == r {
self.tree[u] = arr[l - 1];
return;
}
let mid = (l + r) >> 1;
self.build(u << 1, l, mid, arr);
self.build(u << 1 | 1, mid + 1, r, arr);
self.tree[u] = Self::gcd(self.tree[u << 1], self.tree[u << 1 | 1]);
}

fn modify(&mut self, u: usize, l: usize, r: usize, x: usize, v: i32) {
if l == r {
self.tree[u] = v;
return;
}
let mid = (l + r) >> 1;
if x <= mid {
self.modify(u << 1, l, mid, x, v);
} else {
self.modify(u << 1 | 1, mid + 1, r, x, v);
}
self.tree[u] = Self::gcd(self.tree[u << 1], self.tree[u << 1 | 1]);
}

fn query(&self, u: usize, l: usize, r: usize, ql: usize, qr: usize) -> i32 {
if ql > qr {
return 0;
}
if ql <= l && r <= qr {
return self.tree[u];
}
let mid = (l + r) >> 1;
if qr <= mid {
return self.query(u << 1, l, mid, ql, qr);
}
if ql > mid {
return self.query(u << 1 | 1, mid + 1, r, ql, qr);
}
Self::gcd(
self.query(u << 1, l, mid, ql, mid),
self.query(u << 1 | 1, mid + 1, r, mid + 1, qr),
)
}

fn gcd(mut a: i32, mut b: i32) -> i32 {
while b != 0 {
let t = a % b;
a = b;
b = t;
}
a
}
}

impl Solution {
pub fn count_good_subseq(nums: Vec<i32>, p: i32, queries: Vec<Vec<i32>>) -> i32 {
let n = nums.len();
let mut nums = nums;
let mut seg = SegTree::new(n);

// 初始化线段树：p 的倍数保留原值，否则为 0
let mut init = vec![0; n];
let mut cnt = 0;
for i in 0..n {
if nums[i] % p == 0 {
init[i] = nums[i];
cnt += 1;
}
}
if n > 0 {
seg.build(1, 1, n, &init);
}

let mut ans = 0;
for q in queries {
let idx = q[0] as usize;
let val = q[1];

// 移除旧值
if nums[idx] % p == 0 {
seg.modify(1, 1, n, idx + 1, 0);
cnt -= 1;
}
// 加入新值
if val % p == 0 {
seg.modify(1, 1, n, idx + 1, val);
cnt += 1;
}
nums[idx] = val;

// 整体 GCD 不是 p，不可能有好子序列
if seg.tree[1] != p {
continue;
}

// 不是全部元素都是 p 的倍数，或 n > 6
if cnt < n || n > 6 {
ans += 1;
continue;
}

// n <= 6 且所有元素都是 p 的倍数，暴力枚举删除一个位置
let mut ok = false;
for i in 1..=n {
let left_g = seg.query(1, 1, n, 1, i.saturating_sub(1));
let right_g = seg.query(1, 1, n, i + 1, n);
if SegTree::gcd(left_g, right_g) == p {
ok = true;
break;
}
}
if ok {
ans += 1;
}
}

ans
}
}
```

---

复杂度

- 时间复杂度：O((n + q) \times \log n)，其中 n \le 6 时的暴力枚举是常数级。
- 空间复杂度：O(n)。