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}$.

Python3JavaC++Go

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

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

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

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

Comments