[leetcode 3321][Go][有序集合]计算子数组的 x-sum II

题目链接

leetcode 3321 计算子数组的 x-sum II

前置题目

leetcode 3318 计算子数组的 x-sum I

解题思路

根据题目，首先可以定义一个键值对绑定数值以及该数值在当前子数组中出现的频率 

1
2
3
4
5
//也可以用长度为2的数组，根据个人喜好即可
type pair struct {
	freq int //频率
	num  int //数值
}
然后需要实现以下功能:

  • 根据数值快速查询到子数组频率并修改
  • 修改后的键值对应该保持有序性
  • 能够快速查询子数组前x个频繁出现的数值的和

要实现以上三个功能，可以使用一个滑动窗口加两个有序集合和一个哈希表来实现。

  • 滑动窗口用于表示当前子数组
  • 两个有序集合分别存储当前子数组下前x个频繁出现的数和剩余的数
  • 哈希表用于存储当前子数组下每个数的出现频率

定义比较器

1
2
3
4
5
6
func comparator(a pair, b pair) int {
	if a.freq != b.freq {
		return a.freq - b.freq
	}
	return a.num - b.num
}

定义集合结构

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
type assembly struct {
	k        int
	x        int
	xsum     int64 //当前状态下的xsum
	largeSet *treeset.Set[pair] //表示前x个频繁出现的数
	smallSet *treeset.Set[pair] //表示剩余的数
	freqs    map[int]int //freqs[n]代表当前状态下n出现的频率
}

func New(n int, k int, x int) assembly {
	var assem = assembly{}
	assem.freqs = make(map[int]int)
	assem.largeSet = treeset.NewWith(comparator)
	assem.smallSet = treeset.NewWith(comparator)
	assem.k = k
	assem.x = x
	return assem
}

//因为treeset并未实现Max()和Min()所以这里手动实现
//也可以使用treemap，根据个人爱好选择
//treeset迭代器按照中序遍历的顺序，所以第一个为最小值最后一个为最大值
func findMin[T comparable](set *treeset.Set[T]) (T, bool) {
	t := set.Iterator()
	if t.First() {
		return t.Value(), true
	}
	var zero T
	return zero, false
}

func findMax[T comparable](set *treeset.Set[T]) (T, bool) {
	t := set.Iterator()
	if t.Last() {
		return t.Value(), true
	}
	var zero T
	return zero, false
}

子数组的数据修改操作

当新的元素数值已存在时，应直接更新频率并重新插入有序集合，若不存在则创建新的键值对并插入有序集合。

当子数组数据发生更改后，可能出现以下两种情况:

  • largeSet的元素数量小于x，但smallSet却不为空
  • smallSet中的最大值大于largeSet的最小值（即comparator(findMin(largeSet), findMax(smallSet)) < 0）

此时应该对两个数组进行更新:

  • 若largeSet的元素数量等于x，则弹出largeSet的最小值并加入smallSet，再将smallSet的最大值弹出并加入largeSet
  • 若largeSet的元素数量小于x，则直接弹出smallSet的最大值并加入largeSet

重复以上操作直到largeSet的最小值大于smallSet的最大值或者smallSet为空

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
func (a *assembly) add(n int) { //向子数组中添加元素
	np := pair{num: n, freq: a.freqs[n]}
	if a.largeSet.Contains(np) {
		a.largeSet.Remove(np) //largeSet的修改应该伴随xsum的更新
		a.xsum -= int64(np.num) * int64(np.freq)
	} else {
		a.smallSet.Remove(np)
	}
	a.freqs[n]++
	np.freq++
	a.smallSet.Add(np)
	a.internal()
}

func (a *assembly) reduce(n int) {
	np := pair{num: n, freq: a.freqs[n]}
	if a.largeSet.Contains(np) {
		a.largeSet.Remove(np)
		a.xsum -= int64(np.num) * int64(np.freq)
	} else {
		a.smallSet.Remove(np)
	}
	a.freqs[n]--
	np.freq--
	a.smallSet.Add(np)
	a.internal()
}

func (a *assembly) internal() {
	var largeMin, _ = findMin(a.largeSet)
	var smallMax, ok2 = findMax(a.smallSet)
	for ok2 && (a.largeSet.Size() < a.x || comparator(largeMin, smallMax) < 0) { 
		if a.largeSet.Size() == a.x {
			a.largeSet.Remove(largeMin)
			a.smallSet.Add(largeMin)
			a.xsum -= int64(largeMin.freq) * int64(largeMin.num)
		}
		a.smallSet.Remove(smallMax)
		a.largeSet.Add(smallMax)
		a.xsum += int64(smallMax.num) * int64(smallMax.freq)
		largeMin, _ = findMin(a.largeSet)
		smallMax, ok2 = findMax(a.smallSet)
	}
}

使用滑动窗口进行遍历

定义l和r定义滑动窗口的左右边界，当r向右滑动时则向集合加入新元素，当子数组长度等于k时则记录xsum，当子数组长度大于k时则l向右滑动并从集合减去相应元素

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
func findXSum(nums []int, k int, x int) []int64 {
	var n = len(nums)
	var assem = New(n, k, x)
	var ret []int64

	for l, r := 0, 0; r < n; r++ {
		assem.add(nums[r])
		if r-l+1 > k {
			assem.reduce(nums[l])
			l++
		}
		if r-l+1 >= k {
			ret = append(ret, assem.xsum)
		}
	}

	return ret
}

时间复杂度分析

findMin,findMax,treeset.Add,treeset.Remove的时间复杂度为，每次操作后internal应该只会进行常数次迭代(每次操作只影响一个元素的频率)，每次滑动最多调用一次assem.Add()和assem.Reduce()所以每次滑动的时间复杂度为，综上，总时间复杂度为