class Solution {
public:
vector<int> findXSum(vector<int>& nums, int k, int x) {
}
};
3318. 计算子数组的 x-sum I
给你一个由 n 个整数组成的数组 nums，以及两个整数 k 和 x。
数组的 x-sum 计算按照以下步骤进行：
x 个元素的每次出现。如果两个元素的出现次数相同，则数值 较大 的元素被认为出现次数更多。注意，如果数组中的不同元素少于 x 个，则其 x-sum 是数组的元素总和。
返回一个长度为 n - k + 1 的整数数组 answer，其中 answer[i] 是 子数组 nums[i..i + k - 1] 的 x-sum。
子数组 是数组内的一个连续 非空 的元素序列。
示例 1：
输入：nums = [1,1,2,2,3,4,2,3], k = 6, x = 2
输出：[6,10,12]
解释：
[1, 1, 2, 2, 3, 4]，只保留元素 1 和 2。因此，answer[0] = 1 + 1 + 2 + 2。[1, 2, 2, 3, 4, 2]，只保留元素 2 和 4。因此，answer[1] = 2 + 2 + 2 + 4。注意 4 被保留是因为其数值大于出现其他出现次数相同的元素（3 和 1）。[2, 2, 3, 4, 2, 3]，只保留元素 2 和 3。因此，answer[2] = 2 + 2 + 2 + 3 + 3。示例 2：
输入：nums = [3,8,7,8,7,5], k = 2, x = 2
输出：[11,15,15,15,12]
解释：
由于 k == x，answer[i] 等于子数组 nums[i..i + k - 1] 的总和。
提示：
1 <= n == nums.length <= 501 <= nums[i] <= 501 <= x <= k <= nums.length原站题解
golang 解法, 执行用时: 7 ms, 内存消耗: 4.7 MB, 提交时间: 2024-10-14 09:58:00
import "github.com/emirpasic/gods/v2/trees/redblacktree"
type pair struct{ c, x int } // 出现次数，元素值
func less(p, q pair) int {
return cmp.Or(p.c-q.c, p.x-q.x)
}
func findXSum(nums []int, k, x int) []int {
L := redblacktree.NewWith[pair, struct{}](less)
R := redblacktree.NewWith[pair, struct{}](less)
sumL := 0 // L 的元素和
cnt := map[int]int{}
add := func(x int) {
p := pair{cnt[x], x}
if p.c == 0 {
return
}
if !L.Empty() && less(p, L.Left().Key) > 0 { // p 比 L 中最小的还大
sumL += p.c * p.x
L.Put(p, struct{}{})
} else {
R.Put(p, struct{}{})
}
}
del := func(x int) {
p := pair{cnt[x], x}
if p.c == 0 {
return
}
if _, ok := L.Get(p); ok {
sumL -= p.c * p.x
L.Remove(p)
} else {
R.Remove(p)
}
}
l2r := func() {
p := L.Left().Key
sumL -= p.c * p.x
L.Remove(p)
R.Put(p, struct{}{})
}
r2l := func() {
p := R.Right().Key
sumL += p.c * p.x
R.Remove(p)
L.Put(p, struct{}{})
}
ans := make([]int, len(nums)-k+1)
for r, in := range nums {
// 添加 in
del(in)
cnt[in]++
add(in)
l := r + 1 - k
if l < 0 {
continue
}
// 维护大小
for !R.Empty() && L.Size() < x {
r2l()
}
for L.Size() > x {
l2r()
}
ans[l] = sumL
// 移除 out
out := nums[l]
del(out)
cnt[out]--
add(out)
}
return ans
}
cpp 解法, 执行用时: 28 ms, 内存消耗: 33.6 MB, 提交时间: 2024-10-14 09:57:08
class Solution {
public:
vector<int> findXSum(const vector<int>& nums, int k, int x) {
using pii = pair<int, int>; // 出现次数，元素值
set<pii> L, R;
int sum_l = 0; // L 的元素和
unordered_map<int, int> cnt;
auto add = [&](int x) {
pii p = {cnt[x], x};
if (p.first == 0) {
return;
}
if (!L.empty() && p > *L.begin()) { // p 比 L 中最小的还大
sum_l += p.first * p.second;
L.insert(p);
} else {
R.insert(p);
}
};
auto del = [&](int x) {
pii p = {cnt[x], x};
if (p.first == 0) {
return;
}
auto it = L.find(p);
if (it != L.end()) {
sum_l -= p.first * p.second;
L.erase(it);
} else {
R.erase(p);
}
};
auto l2r = [&]() {
pii p = *L.begin();
sum_l -= p.first * p.second;
L.erase(p);
R.insert(p);
};
auto r2l = [&]() {
pii p = *R.rbegin();
sum_l += p.first * p.second;
R.erase(p);
L.insert(p);
};
vector<int> ans(nums.size() - k + 1);
for (int r = 0; r < nums.size(); r++) {
// 添加 in
int in = nums[r];
del(in);
cnt[in]++;
add(in);
int l = r + 1 - k;
if (l < 0) {
continue;
}
// 维护大小
while (!R.empty() && L.size() < x) {
r2l();
}
while (L.size() > x) {
l2r();
}
ans[l] = sum_l;
// 移除 out
int out = nums[l];
del(out);
cnt[out]--;
add(out);
}
return ans;
}
};
java 解法, 执行用时: 10 ms, 内存消耗: 43.9 MB, 提交时间: 2024-10-14 09:54:42
class Solution {
private final TreeSet<int[]> L = new TreeSet<>((a, b) -> a[0] != b[0] ? a[0] - b[0] : a[1] - b[1]);
private final TreeSet<int[]> R = new TreeSet<>(L.comparator());
private final Map<Integer, Integer> cnt = new HashMap<>();
private int sumL = 0;
public int[] findXSum(int[] nums, int k, int x) {
int[] ans = new int[nums.length - k + 1];
for (int r = 0; r < nums.length; r++) {
// 添加 in
int in = nums[r];
del(in);
cnt.merge(in, 1, Integer::sum); // cnt[in]++
add(in);
int l = r + 1 - k;
if (l < 0) {
continue;
}
// 维护大小
while (!R.isEmpty() && L.size() < x) {
r2l();
}
while (L.size() > x) {
l2r();
}
ans[l] = sumL;
// 移除 out
int out = nums[l];
del(out);
cnt.merge(out, -1, Integer::sum); // cnt[out]--
add(out);
}
return ans;
}
// 添加元素
private void add(int val) {
int c = cnt.get(val);
if (c == 0) {
return;
}
int[] p = new int[]{c, val};
if (!L.isEmpty() && L.comparator().compare(p, L.first()) > 0) { // p 比 L 中最小的还大
sumL += p[0] * p[1];
L.add(p);
} else {
R.add(p);
}
}
// 删除元素
private void del(int val) {
int c = cnt.getOrDefault(val, 0);
if (c == 0) {
return;
}
int[] p = new int[]{c, val};
if (L.contains(p)) {
sumL -= p[0] * p[1];
L.remove(p);
} else {
R.remove(p);
}
}
// 从 L 移动一个元素到 R
private void l2r() {
int[] p = L.pollFirst();
sumL -= p[0] * p[1];
R.add(p);
}
// 从 R 移动一个元素到 L
private void r2l() {
int[] p = R.pollLast();
sumL += p[0] * p[1];
L.add(p);
}
}
python3 解法, 执行用时: 112 ms, 内存消耗: 17.1 MB, 提交时间: 2024-10-14 09:51:35
from sortedcontainers import SortedList
class Solution:
def findXSum(self, nums: List[int], k: int, x: int) -> List[int]:
cnt = defaultdict(int)
L = SortedList() # 保存 tuple (出现次数，元素值)
R = SortedList()
sum_l = 0 # L 的元素和
def add(val: int) -> None:
if cnt[val] == 0:
return
p = (cnt[val], val)
if L and p > L[0]: # p 比 L 中最小的还大
nonlocal sum_l
sum_l += p[0] * p[1]
L.add(p)
else:
R.add(p)
def remove(val: int) -> None:
if cnt[val] == 0:
return
p = (cnt[val], val)
if p in L:
nonlocal sum_l
sum_l -= p[0] * p[1]
L.remove(p)
else:
R.remove(p)
def l2r() -> None:
nonlocal sum_l
p = L[0]
sum_l -= p[0] * p[1]
L.remove(p)
R.add(p)
def r2l() -> None:
nonlocal sum_l
p = R[-1]
sum_l += p[0] * p[1]
R.remove(p)
L.add(p)
ans = [0] * (len(nums) - k + 1)
for r, in_ in enumerate(nums):
# 添加 in_
remove(in_)
cnt[in_] += 1
add(in_)
l = r + 1 - k
if l < 0:
continue
# 维护大小
while R and len(L) < x:
r2l()
while len(L) > x:
l2r()
ans[l] = sum_l
# 移除 out
out = nums[l]
remove(out)
cnt[out] -= 1
add(out)
return ans