LeetCode 2916. 子数组不同元素数目的平方和 IIJava 题解题目描述给你一个整数数组 nums。一个子数组的 不同计数 定义为该子数组中不同元素的个数。请返回 nums 的所有子数组的 不同计数的平方和。答案可能很大对 10^9 7 取模。与第 I 版本的区别I 版求的是“不同计数之和”II 版升级为“平方和”数据范围可达 10^5需要 O(n log n) 解法。---核心思路我们可以枚举子数组的右端点 i并在移动过程中维护一个数组 A其中 A[j] 表示 以 j 开头、以 i 结尾的子数组 的不同元素个数。对于当前右端点 i设 nums[i] 上一次出现的位置为 p若未出现过则为 -1。那么· 对于所有起点 j ∈ [p1, i]新加入的 nums[i] 会为这些子数组增加一个不同的新元素即 A[j] 需要 加 1。· 对于 j ∈ [0, p]nums[i] 已经出现过不同计数不变。因此我们需要一个数据结构支持· 区间加将 [p1, i] 所有元素 1· 全局平方和查询每次更新后求 Σ A[j]^2累加到答案若直接维护数组每次区间加和求和是 O(n)总复杂度 O(n^2) 不可行。我们可以使用 带懒标记的线段树 高效地维护区间和与区间平方和。---线段树维护平方和当区间内每个元素都加上 v 时新平方和 Σ (x v)² Σ (x² 2v·x v²) 旧平方和 2v·Σx v²·len因此我们需要同时维护· sum区间和· sq区间平方和更新公式对长度为 len 的区间加 vjavasq sq 2*v*sum v*v*len;sum sum v*len;所有运算需要对 MOD 1_000_000_007 取模。---算法步骤1. 初始化线段树大小为 n初始值全为 0。2. 用 HashMap 或数组记录每个数字最后出现的位置 last。3. 遍历右端点 i 从 0 到 n-1· 查出 p last.get(nums[i])默认为 -1。· 对线段树区间 [p1, i] 执行加 1 操作。· 将整棵树的平方和即当前所有 j ≤ i 的 A[j]^2 之和累加到答案。· 更新 lastlast.put(nums[i], i)。4. 返回答案取模。复杂度每次区间更新 O(log n)共 n 次故 时间复杂度 O(n log n)空间 O(n)。---Java 代码实现javaclass Solution {private static final int MOD 1_000_000_007;public int sumCounts(int[] nums) {int n nums.length;SegmentTree seg new SegmentTree(n);java.util.MapInteger, Integer last new java.util.HashMap();long ans 0;for (int i 0; i n; i) {int x nums[i];int p last.getOrDefault(x, -1);int ql p 1;int qr i;// 区间 [p1, i] 整体加 1if (ql qr) {seg.update(1, 0, n - 1, ql, qr, 1);}// 累加当前右端点 i 对应的所有子数组平方和ans (ans seg.getTotalSq()) % MOD;last.put(x, i);}return (int) ans;}}class SegmentTree {private final int n;private final long[] sum; // 区间和private final long[] sq; // 区间平方和private final long[] lazy; // 懒标记加数private static final int MOD 1_000_000_007;public SegmentTree(int n) {this.n n;sum new long[4 * n];sq new long[4 * n];lazy new long[4 * n];}// 对节点 idx对应区间 [l, r]整体加 vprivate void apply(int idx, int l, int r, long v) {long len r - l 1;// sq sq 2*v*sum v*v*lensq[idx] (sq[idx] 2L * v % MOD * sum[idx] % MOD v * v % MOD * len % MOD) % MOD;// sum sum v*lensum[idx] (sum[idx] v * len % MOD) % MOD;// 累加懒标记lazy[idx] (lazy[idx] v) % MOD;}// 下推懒标记private void pushDown(int idx, int l, int r) {if (lazy[idx] ! 0 l ! r) {int mid (l r) / 2;apply(idx * 2, l, mid, lazy[idx]);apply(idx * 2 1, mid 1, r, lazy[idx]);lazy[idx] 0;}}// 区间更新对 [ql, qr] 每个元素加 vpublic void update(int idx, int l, int r, int ql, int qr, long v) {if (ql l r qr) {apply(idx, l, r, v);return;}pushDown(idx, l, r);int mid (l r) / 2;if (ql mid) update(idx * 2, l, mid, ql, qr, v);if (qr mid) update(idx * 2 1, mid 1, r, ql, qr, v);sum[idx] (sum[idx * 2] sum[idx * 2 1]) % MOD;sq[idx] (sq[idx * 2] sq[idx * 2 1]) % MOD;}// 获取当前全局平方和public long getTotalSq() {return sq[1];}}---示例验证以 nums [1,2,1] 为例· i01更新 [0,0]平方和 1ans 1· i12更新 [0,1]平方和 1²2²? 实际上 A[2,1]平方和 5ans 156· i21更新 [1,2]A 变为 [2,2,1]平方和 4419ans 6915返回 15符合预期。---总结此题的关键在于将“所有子数组的不同计数平方和”转化为动态维护每个左端点的不同计数并利用线段树区间加结合平方和递推公式高效计算。这种数据结构维护平方和的技巧在多种计数问题中都很实用。