题目描述
给定一个整数数组 nums 和一个整数 k。你可以进行任意次以下操作：
返回确保 至少 有一个大小为 k 的 nums 中的 子数组 的所有元素都相等的所需的 最小 操作数。
示例 1：
输入：nums = [4,-3,2,1,-4,6], k = 3
输出：5
解释：
- 使用 4 次操作来给
nums[1] 增加 4。结果数组为 [4, 1, 2, 1, -4, 6]。
- 使用 1 次操作来给
nums[2] 减少 1。结果数组为 [4, 1, 1, 1, -4, 6]。
- 现在数组包含一个大小为
k = 3 的子数组 [1, 1, 1]，所有元素都想等。因此，答案为 5。
示例 2：
输入：nums = [-2,-2,3,1,4], k = 2
输出：0
解释：
提示：
2 <= nums.length <= 105-106 <= nums[i] <= 1062 <= k <= nums.length
解法
方法一：有序集合
根据题目描述，我们需要找到一个长度为 \(k\) 的子数组，通过最少的操作使得子数组中的所有元素相等，即我们需要找到一个长度为 \(k\) 的子数组，使得子数组中所有元素变成这 \(k\) 个元素的中位数所需的操作次数最小。
我们可以使用两个有序集合 \(l\) 和 \(r\) 分别维护 \(k\) 个元素的左右两部分，其中 \(l\) 用于存储 \(k\) 个元素中较小的一部分，\(r\) 用于存储 \(k\) 个元素中较大的一部分，并且 \(l\) 的元素个数要么等于 \(r\) 的元素个数，要么比 \(r\) 的元素个数少一个，这样 \(r\) 的最小值就是 \(k\) 个元素中的中位数。
时间复杂度 \(O(n \times \log k)\)，空间复杂度 \(O(k)\)。其中 \(n\) 为数组 \(\textit{nums}\) 的长度。
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 | class Solution:
def minOperations(self, nums: List[int], k: int) -> int:
l = SortedList()
r = SortedList()
s1 = s2 = 0
ans = inf
for i, x in enumerate(nums):
l.add(x)
s1 += x
y = l.pop()
s1 -= y
r.add(y)
s2 += y
if len(r) - len(l) > 1:
y = r.pop(0)
s2 -= y
l.add(y)
s1 += y
if i >= k - 1:
ans = min(ans, s2 - r[0] * len(r) + r[0] * len(l) - s1)
j = i - k + 1
if nums[j] in r:
r.remove(nums[j])
s2 -= nums[j]
else:
l.remove(nums[j])
s1 -= nums[j]
return ans
|
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
45
46
47
48
49
50
51
52 | class Solution {
public long minOperations(int[] nums, int k) {
TreeMap<Integer, Integer> l = new TreeMap<>();
TreeMap<Integer, Integer> r = new TreeMap<>();
long s1 = 0, s2 = 0;
int sz1 = 0, sz2 = 0;
long ans = Long.MAX_VALUE;
for (int i = 0; i < nums.length; ++i) {
l.merge(nums[i], 1, Integer::sum);
s1 += nums[i];
++sz1;
int y = l.lastKey();
if (l.merge(y, -1, Integer::sum) == 0) {
l.remove(y);
}
s1 -= y;
--sz1;
r.merge(y, 1, Integer::sum);
s2 += y;
++sz2;
if (sz2 - sz1 > 1) {
y = r.firstKey();
if (r.merge(y, -1, Integer::sum) == 0) {
r.remove(y);
}
s2 -= y;
--sz2;
l.merge(y, 1, Integer::sum);
s1 += y;
++sz1;
}
if (i >= k - 1) {
ans = Math.min(ans, s2 - r.firstKey() * sz2 + r.firstKey() * sz1 - s1);
int j = i - k + 1;
if (r.containsKey(nums[j])) {
if (r.merge(nums[j], -1, Integer::sum) == 0) {
r.remove(nums[j]);
}
s2 -= nums[j];
--sz2;
} else {
if (l.merge(nums[j], -1, Integer::sum) == 0) {
l.remove(nums[j]);
}
s1 -= nums[j];
--sz1;
}
}
}
return ans;
}
}
|
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 | class Solution {
public:
long long minOperations(vector<int>& nums, int k) {
multiset<int> l, r;
long long s1 = 0, s2 = 0, ans = 1e18;
for (int i = 0; i < nums.size(); ++i) {
l.insert(nums[i]);
s1 += nums[i];
int y = *l.rbegin();
l.erase(l.find(y));
s1 -= y;
r.insert(y);
s2 += y;
if (r.size() - l.size() > 1) {
y = *r.begin();
r.erase(r.find(y));
s2 -= y;
l.insert(y);
s1 += y;
}
if (i >= k - 1) {
long long x = *r.begin();
ans = min(ans, s2 - x * (int) r.size() + x * (int) l.size() - s1);
int j = i - k + 1;
if (r.contains(nums[j])) {
r.erase(r.find(nums[j]));
s2 -= nums[j];
} else {
l.erase(l.find(nums[j]));
s1 -= nums[j];
}
}
}
return ans;
}
};
|
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
45
46 | func minOperations(nums []int, k int) int64 {
l := redblacktree.New[int, int]()
r := redblacktree.New[int, int]()
merge := func(st *redblacktree.Tree[int, int], x, v int) {
c, _ := st.Get(x)
if c+v == 0 {
st.Remove(x)
} else {
st.Put(x, c+v)
}
}
var s1, s2, sz1, sz2 int
ans := math.MaxInt64
for i, x := range nums {
merge(l, x, 1)
s1 += x
y := l.Right().Key
merge(l, y, -1)
s1 -= y
merge(r, y, 1)
s2 += y
sz2++
if sz2-sz1 > 1 {
y = r.Left().Key
merge(r, y, -1)
s2 -= y
sz2--
merge(l, y, 1)
s1 += y
sz1++
}
if j := i - k + 1; j >= 0 {
ans = min(ans, s2-r.Left().Key*sz2+r.Left().Key*sz1-s1)
if _, ok := r.Get(nums[j]); ok {
merge(r, nums[j], -1)
s2 -= nums[j]
sz2--
} else {
merge(l, nums[j], -1)
s1 -= nums[j]
sz1--
}
}
}
return int64(ans)
}
|