2025-05-12：计算子数组的 x-sumⅠ。用go语言，给定一个长度为 n 的整数数组 nums，以及两个整数 k 和 x

福大大架构师每日一题

发布于 2025-05-12 09:55:58

10000

代码可运行

文章被收录于专栏：福大大架构师每日一题福大大架构师每日一题

运行总次数：0

代码可运行

2025-05-12：计算子数组的 x-sumⅠ。用go语言，给定一个长度为 n 的整数数组 nums，以及两个整数 k 和 x。

定义数组的 x-sum 如下：

1. 统计数组中各个元素的出现频率。
2. 选出出现次数最多的前 x 个元素的所有出现位置。若出现次数相同，则数值较大的元素优先被选中。
3. 将选中的这些元素加起来，得到 x-sum。

如果不同的元素数量少于 x，则直接返回数组所有元素的和。

请你计算数组中所有长度为 k 的连续子数组的 x-sum，返回一个长度为 n - k + 1 的数组 answer，其中 answer[i] 表示子数组 nums[i..i + k - 1] 的 x-sum。

1 <= n == nums.length <= 50。

1 <= nums[i] <= 50。

1 <= x <= k <= nums.length。

输入：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。

目来自leetcode3318。

解决步骤

1. 滑动窗口：我们需要遍历所有长度为 k 的子数组。滑动窗口的起始位置从 0 到 n - k。
2. 频率统计：对于每个子数组，统计每个元素的出现频率。
3. 选择前 x 个元素：
- • 将频率和元素值组合成 (频率, 元素值) 的二元组。
- • 对这些二元组排序：先按频率降序，频率相同则按元素值降序。
- • 选择前 x 个二元组对应的元素值。
4. 计算 x-sum：
- • 遍历子数组，累加被选中的元素的值。
5. 优化：
- • 使用红黑树（L 和 R）来维护频率和元素值的排序。
- • L 存储前 x 个最高频率的元素，R 存储剩余元素。
- • sumL 记录 L 中所有元素的频率乘以元素值的和（即 x-sum）。
- • 动态维护 L 和 R 的大小关系：确保 L 的大小为 x。

详细过程

1. 初始化：
- • L 和 R 是两个红黑树，用于存储 (频率, 元素值) 的二元组。
- • sumL 记录 L 中所有元素的 频率 * 元素值 的和。
- • cnt 是一个哈希表，记录当前窗口中每个元素的出现频率。
2. 滑动窗口：
- • 遍历数组，每次移动窗口时：
  - • 移除左边界的元素（如果窗口已满）。
  - • 添加右边界的元素。
  - • 更新 cnt 和红黑树。
3. 维护红黑树：
- • 添加或删除元素时，根据 (频率, 元素值) 更新 L 或 R。
- • 确保 L 的大小为 x：
  - • 如果 L 的大小小于 x，从 R 中移动最大的元素到 L。
  - • 如果 L 的大小大于 x，从 L 中移动最小的元素到 R。
4. 计算 x-sum：
- • sumL 即为当前窗口的 x-sum。

时间复杂度

• 滑动窗口遍历：O(n)。
• 每次窗口移动：
- • 更新 cnt：O(1)。
- • 红黑树操作（插入、删除、查找）：O(log m)，其中 m 是窗口中不同元素的数量（最多 50）。
- • 维护 L 和 R 的大小：最多 O(log m) 操作。
• 总时间复杂度：O(n * log m)，其中 m <= 50，因此可以认为是 O(n)。

空间复杂度

• cnt：O(m)，m 是不同元素的数量（最多 50）。
• L 和 R：O(m)。
• 总空间复杂度：O(m)，即 O(1)（因为 m 最多为 50）。

最终答案

总时间复杂度：O(n)（因为 log m 是常数）。总额外空间复杂度：O(1)（因为 m 是常数）。

Go完整代码如下：

package main

import (
    "cmp"
    "fmt"
    "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) []int64 {
    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([]int64, 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] = int64(sumL)

        // 移除 out
        out := nums[l]
        del(out)
        cnt[out]--
        add(out)
    }
    return ans
}

func main() {
    nums := []int{1, 1, 2, 2, 3, 4, 2, 3}
    k := 6
    x := 2
    result := findXSum(nums, k, x)
    fmt.Println(result)
}