Skip to content

Latest commit

 

History

History
269 lines (230 loc) · 7.93 KB

README_EN.md

File metadata and controls

269 lines (230 loc) · 7.93 KB
comments difficulty edit_url
true
Medium

3422. Minimum Operations to Make Subarray Elements Equal 🔒

中文文档

Description

You are given an integer array nums and an integer k. You can perform the following operation any number of times:

  • Increase or decrease any element of nums by 1.

Return the minimum number of operations required to ensure that at least one subarray of size k in nums has all elements equal.

 

Example 1:

Input: nums = [4,-3,2,1,-4,6], k = 3

Output: 5

Explanation:

  • Use 4 operations to add 4 to nums[1]. The resulting array is [4, 1, 2, 1, -4, 6].
  • Use 1 operation to subtract 1 from nums[2]. The resulting array is [4, 1, 1, 1, -4, 6].
  • The array now contains a subarray [1, 1, 1] of size k = 3 with all elements equal. Hence, the answer is 5.

Example 2:

Input: nums = [-2,-2,3,1,4], k = 2

Output: 0

Explanation:

  • The subarray [-2, -2] of size k = 2 already contains all equal elements, so no operations are needed. Hence, the answer is 0.

 

Constraints:

  • 2 <= nums.length <= 105
  • -106 <= nums[i] <= 106
  • 2 <= k <= nums.length

Solutions

Solution 1: Ordered Set

According to the problem description, we need to find a subarray of length $k$ and make all elements in the subarray equal with the minimum number of operations. That is, we need to find a subarray of length $k$ such that the minimum number of operations required to make all elements in the subarray equal to the median of these $k$ elements is minimized.

We can use two ordered sets $l$ and $r$ to maintain the left and right parts of the $k$ elements, respectively. $l$ is used to store the smaller part of the $k$ elements, and $r$ is used to store the larger part of the $k$ elements. The number of elements in $l$ is either equal to the number of elements in $r$ or one less than the number of elements in $r$, so the minimum value in $r$ is the median of the $k$ elements.

The time complexity is $O(n \times \log k)$, and the space complexity is $O(k)$. Here, $n$ is the length of the array $\textit{nums}$.

Python3

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

Java

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;
    }
}

C++

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;
    }
};

Go

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)
}