LeetCode 3321. 计算子数组的 x-sum II — Python3 实现

题目概述

给定数组 `nums`，整数 `k` 和 `x`。需要计算每个长度为 `k` 的滑动窗口的 x-sum：

1. 统计窗口中每个元素的出现次数
2. 保留出现次数最多的前 `x` 个元素（出现次数相同则数值大的优先）
3. 计算保留元素的总和

返回长度为 `n - k + 1` 的结果数组。

约束：`1 <= n <= 10^5`，需要 `O(n log k)` 的解法。

---

核心思路：哈希表 + 两个有序集合（SortedList）

使用 滑动窗口 配合两个有序集合：

数据结构    作用    
`l`（Top 集合）    存储当前窗口中出现次数最多的前 `x` 个元素    
`r`（剩余集合）    存储其余元素    
`cnt`（Counter）    记录窗口中每个元素的频次    
`s`    维护 `l` 中所有元素的 `频次 × 数值` 之和    

排序规则：按 `(频次, 数值)` 升序排列。这样 `l[0]` 是 `l` 中最小的（最应该被淘汰的），`r.pop()` 取到的是 `r` 中最大的（最应该晋升的）。

滑动窗口操作：
1. 右端元素进入窗口：先 `remove` 旧状态，更新 `cnt`，再 `add` 新状态
2. 平衡 `l` 和 `r`：确保 `len(l) == x`（或窗口中不同元素不足 `x` 个时取全部）
3. 记录当前窗口的 `x-sum`
4. 左端元素离开窗口：同样先 `remove` 旧状态，更新 `cnt`，再 `add` 新状态

---

Python3 代码

```python
from sortedcontainers import SortedList
from collections import Counter
from typing import List

class Solution:
    def findXSum(self, nums: List[int], k: int, x: int) -> List[int]:
        """
        计算每个长度为 k 的滑动窗口的 x-sum
        
        核心思路：哈希表 + 两个有序集合（SortedList）
        - l: 存储前 x 个高频元素（按频次升序，频次相同按数值升序）
        - r: 存储其余元素
        - s: 维护 l 中所有元素的 (频次 × 数值) 之和
        """
        
        def add(v: int):
            """将元素 v 的当前状态 (cnt[v], v) 加入合适的集合"""
            if cnt[v] == 0:
                return
            p = (cnt[v], v)
            # 如果 l 非空且 v 的状态比 l 中最小的还要大，则加入 l
            if l and p > l[0]:
                nonlocal s
                s += p[0] * p[1]
                l.add(p)
            else:
                r.add(p)
        
        def remove(v: int):
            """从集合中移除元素 v 的当前状态 (cnt[v], v)"""
            if cnt[v] == 0:
                return
            p = (cnt[v], v)
            if p in l:
                nonlocal s
                s -= p[0] * p[1]
                l.remove(p)
            else:
                r.remove(p)
        
        # l: 前x个高频元素的集合（按频次升序，频次相同按数值升序）
        l = SortedList()
        # r: 剩余元素的集合
        r = SortedList()
        # cnt: 记录窗口中每个元素的频次
        cnt = Counter()
        # s: l 集合中所有元素的 (频次 × 数值) 之和
        s = 0
        n = len(nums)
        ans = [0] * (n - k + 1)
        
        for i, v in enumerate(nums):
            # 1. 将当前元素加入窗口
            remove(v)       # 移除旧状态
            cnt[v] += 1     # 频次+1
            add(v)          # 加入新状态
            
            j = i - k + 1   # 窗口左边界
            if j < 0:
                continue    # 窗口还未形成
            
            # 2. 平衡 l 和 r，确保 l 中恰好有 x 个元素
            # 如果 l 中元素不足 x 个，从 r 中补充最大的
            while r and len(l) < x:
                p = r.pop()         # r 中最大的元素
                l.add(p)
                s += p[0] * p[1]    # 加入 l，更新总和
            
            # 如果 l 中元素超过 x 个，将最小的移出到 r
            while len(l) > x:
                p = l.pop(0)        # l 中最小的元素
                s -= p[0] * p[1]    # 从 l 移除，更新总和
                r.add(p)
            
            # 3. 记录当前窗口的 x-sum
            ans[j] = s
            
            # 4. 窗口左边界元素移出
            remove(nums[j])         # 移除旧状态
            cnt[nums[j]] -= 1       # 频次-1
            add(nums[j])            # 加入新状态（频次可能为0，add中会处理）
        
        return ans
```

---

复杂度分析

指标    复杂度    说明    
时间    `O(n × log k)`    每个元素进出窗口各一次，每次 SortedList 操作 `O(log k)`    
空间    `O(k)`    Counter 和两个有序集合最多存储窗口中的 `k` 个不同元素    

---

关键点说明

1. SortedList 的排序：默认按元组 `(频次, 数值)` 的字典序升序排列。`l[0]` 是 `l` 中最小的（最弱），`r.pop()` 取到 `r` 中最大的（最强）。

2. 先 remove 再 add 的原因：元素频次变化时，它在有序集合中的位置可能改变，必须先移除旧状态 `(旧频次, v)`，再插入新状态 `(新频次, v)`。

3. 边界情况：当窗口中不同元素个数 `< x` 时，`l` 会包含所有不同元素，`r` 为空，此时 x-sum 等于窗口总和，符合题意。

4. 关于 `sortedcontainers`：LeetCode 的 Python3 运行环境默认支持 `sortedcontainers` 库，无需额外安装。本地测试时需要 `pip install sortedcontainers`。

---

验证结果

```
示例 1: nums=[1,1,2,2,3,4,2,3], k=6, x=2
结果: [6, 10, 12]  ✓

示例 2: nums=[3,8,7,8,7,5], k=2, x=2
结果: [11, 15, 15, 15, 12]  ✓
```

下载完整代码：[leetcode_3321_solution.py](sandbox:///mnt/agents/output/leetcode_3321_solution.py)