DeepSeek LeetCode 2916. 子数组不同元素数目的平方和 II Python3实现

LeetCode 2916. 子数组不同元素数目的平方和 IIPython 题解题目描述给你一个整数数组 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 的区间加 vpythonsq sq 2*v*sum v*v*lensum sum v*len所有运算需要对 MOD 1_000_000_007 取模。---算法步骤1. 初始化线段树大小为 n初始值全为 0。2. 用字典记录每个数字最后出现的位置 last。3. 遍历右端点 i 从 0 到 n-1· 查出 p last.get(nums[i], -1)。· 对线段树区间 [p1, i] 执行加 1 操作。· 将整棵树的平方和即当前所有 j ≤ i 的 A[j]^2 之和累加到答案。· 更新 last[nums[i]] i。4. 返回答案取模。复杂度每次区间更新 O(log n)共 n 次故 时间复杂度 O(n log n)空间 O(n)。---Python3 代码实现pythonclass SegmentTree:def __init__(self, n: int):self.n nself.MOD 1_000_000_007self.sum [0] * (4 * n)self.sq [0] * (4 * n)self.lazy [0] * (4 * n)def _apply(self, idx: int, l: int, r: int, v: int) - None:对节点 idx 对应区间 [l, r] 整体加 vlength r - l 1v_mod v % self.MOD# sq sq 2*v*sum v*v*lenself.sq[idx] (self.sq[idx] 2 * v_mod * self.sum[idx] v_mod * v_mod * length) % self.MOD# sum sum v*lenself.sum[idx] (self.sum[idx] v_mod * length) % self.MOD# 累加懒标记self.lazy[idx] (self.lazy[idx] v_mod) % self.MODdef _push_down(self, idx: int, l: int, r: int) - None:下推懒标记到子节点if self.lazy[idx] ! 0 and l ! r:mid (l r) // 2self._apply(idx * 2, l, mid, self.lazy[idx])self._apply(idx * 2 1, mid 1, r, self.lazy[idx])self.lazy[idx] 0def update(self, idx: int, l: int, r: int, ql: int, qr: int, v: int) - None:区间 [ql, qr] 每个元素加 vif ql l and r qr:self._apply(idx, l, r, v)returnself._push_down(idx, l, r)mid (l r) // 2if ql mid:self.update(idx * 2, l, mid, ql, qr, v)if qr mid:self.update(idx * 2 1, mid 1, r, ql, qr, v)self.sum[idx] (self.sum[idx * 2] self.sum[idx * 2 1]) % self.MODself.sq[idx] (self.sq[idx * 2] self.sq[idx * 2 1]) % self.MODdef get_total_sq(self) - int:获取当前全局平方和return self.sq[1]class Solution:def sumCounts(self, nums: list[int]) - int:MOD 1_000_000_007n len(nums)seg SegmentTree(n)last {} # 记录每个数字上一次出现的下标ans 0for i, x in enumerate(nums):p last.get(x, -1)ql, qr p 1, iif ql qr:seg.update(1, 0, n - 1, ql, qr, 1)ans (ans seg.get_total_sq()) % MODlast[x] ireturn ans---示例验证以 nums [1,2,1] 为例· i01更新 [0,0]平方和 1ans 1· i12更新 [0,1]A[2,1]平方和 415ans 6· i21更新 [1,2]A[2,2,1]平方和 4419ans 15返回 15符合预期。---总结此题的关键在于将“所有子数组的不同计数平方和”转化为动态维护每个左端点的不同计数并利用线段树区间加结合平方和递推公式高效计算。这种数据结构维护平方和的技巧在多种计数问题中都很实用。