LeetCode Solutions Blog

Minimum Operations to Make Subarray Elements Equal

Number: 3698

Difficulty: Medium

Paid? Yes

Companies: N/A

Problem Description

Given an integer array nums and an integer k, you can increase or decrease any element by 1 at each step. Return the minimum number of operations required so that there exists at least one contiguous subarray of size k in which all elements are equal.

Key Insights

To equalize a subarray using the minimum number of operations, transforming all elements to its median minimizes the sum of absolute differences.
A sliding window approach allows us to efficiently consider every subarray of size k.
Maintaining the median dynamically for each window is essential; this can be achieved using two heaps (a max-heap for the lower half and a min-heap for the upper half).
Update the sliding window by adding the new element and removing the element that goes out of the window.
Aggregated sum information for both heaps is maintained so that the cost to convert the window can be computed in constant time after each update.

Space and Time Complexity

Time Complexity: O(n log k) due to insertion and deletion operations (each O(log k)) for every sliding window update.
Space Complexity: O(k) for storing elements in the two heaps.

Solution

We use a sliding window of size k and maintain its median using two heaps. The max-heap (using negative values in Python) stores the lower half and the min-heap stores the upper half of the window elements. For each window, the cost (total operations) to convert all elements to the median is computed by: cost = median * (size of max-heap) - (sum of max-heap) + (sum of min-heap) - median * (size of min-heap).

As the window slides, we update our heaps (adding the new element and removing the exiting element) and recalculate the cost, tracking the minimum cost encountered.

Code Solutions

# Python solution using two heaps
import heapq

class MedianMaintenance:
    def __init__(self):
        # max_heap stores negative values to simulate a max-heap; min_heap is a normal min-heap.
        self.max_heap = []  # lower half
        self.min_heap = []  # upper half
        self.sum_max = 0    # sum of elements in max_heap (as positive values)
        self.sum_min = 0    # sum of elements in min_heap

    def add_num(self, num):
        if not self.max_heap or num <= -self.max_heap[0]:
            heapq.heappush(self.max_heap, -num)
            self.sum_max += num
        else:
            heapq.heappush(self.min_heap, num)
            self.sum_min += num
        self.balance_heaps()

    def remove_num(self, num):
        # Remove num from the appropriate heap and rebalance.
        if self.max_heap and num <= -self.max_heap[0]:
            self.max_heap.remove(-num)
            heapq.heapify(self.max_heap)
            self.sum_max -= num
        else:
            self.min_heap.remove(num)
            heapq.heapify(self.min_heap)
            self.sum_min -= num
        self.balance_heaps()

    def balance_heaps(self):
        # Ensure the heaps' sizes differ by at most 1.
        if len(self.max_heap) > len(self.min_heap) + 1:
            moved = -heapq.heappop(self.max_heap)
            self.sum_max -= moved
            heapq.heappush(self.min_heap, moved)
            self.sum_min += moved
        elif len(self.min_heap) > len(self.max_heap):
            moved = heapq.heappop(self.min_heap)
            self.sum_min -= moved
            heapq.heappush(self.max_heap, -moved)
            self.sum_max += moved

    def get_median(self):
        return -self.max_heap[0]

    def get_cost(self):
        median = self.get_median()
        size_max = len(self.max_heap)
        size_min = len(self.min_heap)
        return median * size_max - self.sum_max + self.sum_min - median * size_min

def min_operations(nums, k):
    n = len(nums)
    median_maint = MedianMaintenance()
    # Initialize the first window.
    for i in range(k):
        median_maint.add_num(nums[i])
    min_cost = median_maint.get_cost()
    # Slide the window.
    for i in range(k, n):
        median_maint.add_num(nums[i])
        median_maint.remove_num(nums[i-k])
        min_cost = min(min_cost, median_maint.get_cost())
    return min_cost

# Example usage:
nums = [4, -3, 2, 1, -4, 6]
k = 3
print(min_operations(nums, k))  # Expected output: 5