Kimi LeetCode 3569. 分割数组后不同质数的最大数目 Java实现
以下是 LeetCode 3569 分割数组后不同质数的最大数目 的 Java 实现。
核心思路
问题转化
对于每个查询，我们需要找到最优分割点 k，使得：
\text{prefix中不同质数个数} + \text{suffix中不同质数个数}
关键观察：设整个数组中共有 T 个不同的质数。若某个质数 p 的首次出现位置为 l、末次出现位置为 r，则：
- 若分割点 k \in [l+1, r]，则 p 在前后两段都出现，贡献为 2（额外收益 +1）
- 否则 p 只出现在一段中，贡献为 1
因此：
\text{答案} = T + \max{k} \{\text{在位置 } k \text{ 处被"两边都有"的质数个数}\}
数据结构
- 埃氏筛：O(10^5) 预处理质数表，O(1) 判断质数
- `TreeSet`（红黑树）：对每个质数维护其出现的所有下标，支持 O(\log n) 的插入、删除、查最值
- 懒标记线段树：维护每个分割位置的"额外收益"（即有多少质数在该位置被两边都包含），支持 O(\log n) 的区间加减和区间最大值查询
每次查询修改一个元素时，只需更新受影响质数的区间贡献，然后查询线段树最大值即可。
时间复杂度 O((n+q) \cdot \log n)，空间复杂度 O(n)。
```java
import java.util.*;
class Solution {
static final int MX = 100000;
static boolean[] isPrime = new boolean[MX + 1];
static boolean sieveInitialized = false;
// 懒标记线段树：区间加 + 区间最大值
static class SegTree {
int n, size;
int[] mx, lazy;
SegTree(int n) {
this.n = n;
this.size = 1;
while (size < n) size <<= 1;
mx = new int[size << 1];
lazy = new int[size << 1];
}
void apply(int o, int v) {
mx[o] += v;
lazy[o] += v;
}
void push(int o) {
if (lazy[o] != 0) {
apply(o << 1, lazy[o]);
apply(o << 1 | 1, lazy[o]);
lazy[o] = 0;
}
}
void update(int l, int r, int v) {
if (l > r) return;
update(1, 0, size - 1, l, r, v);
}
void update(int o, int L, int R, int l, int r, int v) {
if (l <= L && R <= r) {
apply(o, v);
return;
}
push(o);
int mid = (L + R) >> 1;
if (l <= mid) update(o << 1, L, mid, l, r, v);
if (r > mid) update(o << 1 | 1, mid + 1, R, l, r, v);
mx[o] = Math.max(mx[o << 1], mx[o << 1 | 1]);
}
int queryMax() {
return mx[1];
}
}
public int[] maximumCount(int[] nums, int[][] queries) {
// 埃氏筛预处理质数表（只执行一次）
if (!sieveInitialized) {
Arrays.fill(isPrime, true);
isPrime[0] = isPrime[1] = false;
for (int i = 2; i * i <= MX; i++) {
if (isPrime[i]) {
for (int j = i * i; j <= MX; j += i) {
isPrime[j] = false;
}
}
}
sieveInitialized = true;
}
int n = nums.length;
int q = queries.length;
int[] ans = new int[q];
// pos.get(p) = 质数 p 出现的所有下标（有序）
Map<Integer, TreeSet<Integer>> pos = new HashMap<>();
for (int i = 0; i < n; i++) {
int x = nums[i];
if (isPrime[x]) {
pos.computeIfAbsent(x, k -> new TreeSet<>()).add(i);
}
}
// 线段树：维护每个分割位置的"额外收益"
SegTree seg = new SegTree(n);
// 初始化：出现 ≥2 次的质数，其 [first, last] 区间贡献 +1
for (TreeSet<Integer> indices : pos.values()) {
if (indices.size() >= 2) {
seg.update(indices.first(), indices.last(), 1);
}
}
// 处理每个查询
for (int qi = 0; qi < q; qi++) {
int idx = queries[qi][0];
int val = queries[qi][1];
int old = nums[idx];
nums[idx] = val;
// 删除旧值
if (isPrime[old]) {
TreeSet<Integer> indices = pos.get(old);
if (indices.size() >= 2) {
seg.update(indices.first(), indices.last(), -1);
}
indices.remove(idx);
if (indices.isEmpty()) {
pos.remove(old);
} else if (indices.size() >= 2) {
seg.update(indices.first(), indices.last(), 1);
}
}
// 插入新值
if (isPrime[val]) {
TreeSet<Integer> indices = pos.get(val);
if (indices == null) {
indices = new TreeSet<>();
pos.put(val, indices);
} else {
if (indices.size() >= 2) {
seg.update(indices.first(), indices.last(), -1);
}
}
indices.add(idx);
if (indices.size() >= 2) {
seg.update(indices.first(), indices.last(), 1);
}
}
// 答案 = 总不同质数个数 + 最大额外收益
ans[qi] = pos.size() + seg.queryMax();
}
return ans;
}
}
```
验证示例
输入 输出 说明
`nums = [2,1,3,1,2]`, `queries = [[1,2],[3,3]]` `[3, 4]` 第一次查询后分割得 1+2=3；第二次得 2+2=4
`nums = [2,1,4]`, `queries = [[0,1]]` `[0]` 修改后数组无质数，答案为 0