LeetCode 3505 - Minimum Operations to Make Elements Within K Subarrays Equal

You are given an integer array nums, and you want to create at least k non-overlapping subarrays. Every chosen subarray must have length exactly x, and after performing some modifications, every element inside each chosen subarray must become equal.

LeetCode Problem 3505

Difficulty: 🔴 Hard
Topics: Array, Hash Table, Math, Dynamic Programming, Sliding Window, Heap (Priority Queue)

Solution

Problem Understanding

You are given an integer array nums, and you want to create at least k non-overlapping subarrays. Every chosen subarray must have length exactly x, and after performing some modifications, every element inside each chosen subarray must become equal.

The only allowed operation is increasing or decreasing a single array element by 1. Every unit change costs exactly one operation. If an element changes from 5 to 1, that contributes 4 operations. If an element changes from -2 to 3, that contributes 5 operations.

The goal is to minimize the total number of operations required so that there exist at least k non-overlapping length-x subarrays whose elements are all equal.

A very important observation is that different chosen subarrays do not overlap. Because they are disjoint, the cost of transforming one chosen subarray is independent of every other chosen subarray. This allows us to compute the optimal cost for every possible window of length x, then choose the best set of k non-overlapping windows.

The constraints are large:

  • n = nums.length can be up to 100000
  • x can also be as large as 100000
  • k <= 15

The small value of k is a strong hint that dynamic programming over the number of selected subarrays is likely possible. The large value of n means we must compute window costs much more efficiently than sorting every window independently.

Some important edge cases include:

  • Windows that already contain equal values, whose cost is 0.
  • Negative values and large positive values.
  • Cases where many candidate windows overlap heavily.
  • Cases where x is very large, close to n.
  • Cases where multiple optimal answers exist, where we only need the minimum cost.

Approaches

Brute Force

A straightforward approach would be:

  1. Enumerate every length-x window.
  2. Sort the elements of the window.
  3. Compute the minimum cost to make all elements equal.
  4. Use dynamic programming to choose k non-overlapping windows.

The minimum cost for a window is obtained by converting every element to the median of that window.

However, there are O(n) windows, and sorting each window costs O(x log x). The total complexity becomes:

O(n * x log x)

With both n and x potentially reaching 100000, this is far too slow.

Key Insight

For any fixed window, the minimum cost of making all values equal is achieved by converting everything to the median.

This is the classic property of absolute deviations:

$\sum |a_i-m|$

The expression is minimized when m is a median of the values.

Therefore, the problem becomes:

  1. Efficiently compute the median-based cost for every sliding window of length x.
  2. Choose k non-overlapping windows with minimum total cost.

The first part can be solved using a sliding-window median structure based on two heaps and lazy deletion.

The second part becomes a small dynamic programming problem because k <= 15.

Approach Time Complexity Space Complexity Notes
Brute Force O(n * x log x) O(x) Sort every window independently
Optimal O(n log x + nk) O(n + kx) Sliding-window median costs plus DP

Algorithm Walkthrough

Step 1: Compute the cost of one window using its median

For a window:

[a1, a2, a3, ..., ax]

the minimum operations needed to make every element equal is:

sum(abs(ai - median))

So we need a data structure that can:

  • Insert a value
  • Remove a value
  • Query the current median
  • Query the total absolute distance to the median

all in logarithmic time.

Step 2: Maintain two heaps

We use the standard sliding-window median structure:

  • small = max heap containing the lower half
  • large = min heap containing the upper half

The median is always the top of small.

We also maintain:

  • small_sum
  • large_sum

These store the sums of elements currently contained in each heap.

Step 3: Compute the window cost in O(1)

Let:

  • m = median
  • L = number of elements in small
  • R = number of elements in large

Then:

cost_left  = m * L - small_sum
cost_right = large_sum - m * R

Therefore:

cost = cost_left + cost_right

or equivalently:

$\text{cost}=m\cdot L-\text{smallSum}+\text{largeSum}-m\cdot R$

This formula gives the sum of absolute distances to the median.

Step 4: Slide the window

Initialize the first window of length x.

Then repeatedly:

  1. Remove the outgoing value.
  2. Insert the incoming value.
  3. Rebalance the heaps.
  4. Compute the new window cost.

This gives:

cost[start]

for every window starting position.

Step 5: Dynamic programming over non-overlapping windows

Let:

dp[t][i]

be the minimum cost after processing the first i array positions and choosing exactly t valid windows.

Initially:

dp[0][0] = 0

For every position i:

Skip position

We may simply move forward:

dp[t][i + 1]

Take window starting at i

If a length-x window fits:

window = [i, i + x - 1]

then:

dp[t + 1][i + x]
    =
min(
    dp[t + 1][i + x],
    dp[t][i] + cost[i]
)

Notice that the next state jumps to i + x, which automatically guarantees non-overlapping windows.

Why it works

For each window, converting all elements to the median yields the minimum possible modification cost. Because chosen windows are non-overlapping, their costs are completely independent and can simply be added together.

The dynamic programming step explores every valid combination of non-overlapping windows. Every state represents all possibilities using the processed prefix of the array, and every transition either skips a position or selects the window beginning there. Since all valid solutions can be represented through these transitions, the DP finds the global minimum.

Python Solution

from typing import List
from collections import defaultdict
import heapq

class DualHeap:
    def __init__(self):
        self.small = []  # max heap via negatives
        self.large = []  # min heap

        self.delayed = defaultdict(int)

        self.small_size = 0
        self.large_size = 0

        self.small_sum = 0
        self.large_sum = 0

    def prune_small(self) -> None:
        while self.small:
            num = -self.small[0]
            if self.delayed[num]:
                self.delayed[num] -= 1
                if self.delayed[num] == 0:
                    del self.delayed[num]
                heapq.heappop(self.small)
            else:
                break

    def prune_large(self) -> None:
        while self.large:
            num = self.large[0]
            if self.delayed[num]:
                self.delayed[num] -= 1
                if self.delayed[num] == 0:
                    del self.delayed[num]
                heapq.heappop(self.large)
            else:
                break

    def rebalance(self) -> None:
        if self.small_size > self.large_size + 1:
            val = -heapq.heappop(self.small)

            self.small_sum -= val
            self.small_size -= 1

            heapq.heappush(self.large, val)
            self.large_sum += val
            self.large_size += 1

            self.prune_small()

        elif self.small_size < self.large_size:
            val = heapq.heappop(self.large)

            self.large_sum -= val
            self.large_size -= 1

            heapq.heappush(self.small, -val)
            self.small_sum += val
            self.small_size += 1

            self.prune_large()

    def add(self, num: int) -> None:
        if not self.small or num <= -self.small[0]:
            heapq.heappush(self.small, -num)
            self.small_sum += num
            self.small_size += 1
        else:
            heapq.heappush(self.large, num)
            self.large_sum += num
            self.large_size += 1

        self.rebalance()

    def remove(self, num: int) -> None:
        self.delayed[num] += 1

        if num <= -self.small[0]:
            self.small_size -= 1
            self.small_sum -= num

            if num == -self.small[0]:
                self.prune_small()
        else:
            self.large_size -= 1
            self.large_sum -= num

            if self.large and num == self.large[0]:
                self.prune_large()

        self.rebalance()

    def cost(self) -> int:
        median = -self.small[0]

        left_cost = median * self.small_size - self.small_sum
        right_cost = self.large_sum - median * self.large_size

        return left_cost + right_cost

class Solution:
    def minOperations(self, nums: List[int], x: int, k: int) -> int:
        n = len(nums)

        dh = DualHeap()

        for i in range(x):
            dh.add(nums[i])

        window_costs = [dh.cost()]

        for right in range(x, n):
            dh.remove(nums[right - x])
            dh.add(nums[right])

            window_costs.append(dh.cost())

        INF = 10**18

        dp = [[INF] * (n + 1) for _ in range(k + 1)]
        dp[0][0] = 0

        for i in range(n):
            for used in range(k + 1):
                if dp[used][i] == INF:
                    continue

                if dp[used][i] < dp[used][i + 1]:
                    dp[used][i + 1] = dp[used][i]

                if used < k and i + x <= n:
                    candidate = dp[used][i] + window_costs[i]
                    if candidate < dp[used + 1][i + x]:
                        dp[used + 1][i + x] = candidate

        return min(dp[k])

Implementation Explanation

The DualHeap class maintains the median of the current window and also tracks the sum of values on each side of the median. This allows the window cost to be computed in constant time once the heaps are balanced.

The sliding window phase builds window_costs, where window_costs[i] is the minimum number of operations needed to make the window starting at index i consist entirely of equal values.

After that, the problem becomes selecting k non-overlapping windows. The DP table uses array positions rather than window indices. When a window is selected, the state jumps directly from i to i + x, guaranteeing that future selections cannot overlap with the chosen window.

The answer is the minimum value among all states that have selected exactly k windows.

Go Solution

package main

import (
	"container/heap"
)

type MaxHeap []int

func (h MaxHeap) Len() int            { return len(h) }
func (h MaxHeap) Less(i, j int) bool  { return h[i] > h[j] }
func (h MaxHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
func (h *MaxHeap) Push(x interface{}) { *h = append(*h, x.(int)) }
func (h *MaxHeap) Pop() interface{} {
	old := *h
	n := len(old)
	v := old[n-1]
	*h = old[:n-1]
	return v
}

type MinHeap []int

func (h MinHeap) Len() int            { return len(h) }
func (h MinHeap) Less(i, j int) bool  { return h[i] < h[j] }
func (h MinHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
func (h *MinHeap) Push(x interface{}) { *h = append(*h, x.(int)) }
func (h *MinHeap) Pop() interface{} {
	old := *h
	n := len(old)
	v := old[n-1]
	*h = old[:n-1]
	return v
}

type DualHeap struct {
	small MaxHeap
	large MinHeap

	delayed map[int]int

	smallSize int
	largeSize int

	smallSum int64
	largeSum int64
}

func NewDualHeap() *DualHeap {
	d := &DualHeap{
		delayed: make(map[int]int),
	}
	heap.Init(&d.small)
	heap.Init(&d.large)
	return d
}

func (d *DualHeap) pruneSmall() {
	for d.small.Len() > 0 {
		num := d.small[0]
		if cnt, ok := d.delayed[num]; ok && cnt > 0 {
			d.delayed[num]--
			if d.delayed[num] == 0 {
				delete(d.delayed, num)
			}
			heap.Pop(&d.small)
		} else {
			break
		}
	}
}

func (d *DualHeap) pruneLarge() {
	for d.large.Len() > 0 {
		num := d.large[0]
		if cnt, ok := d.delayed[num]; ok && cnt > 0 {
			d.delayed[num]--
			if d.delayed[num] == 0 {
				delete(d.delayed, num)
			}
			heap.Pop(&d.large)
		} else {
			break
		}
	}
}

func (d *DualHeap) rebalance() {
	if d.smallSize > d.largeSize+1 {
		val := heap.Pop(&d.small).(int)

		d.smallSize--
		d.smallSum -= int64(val)

		heap.Push(&d.large, val)
		d.largeSize++
		d.largeSum += int64(val)

		d.pruneSmall()
	} else if d.smallSize < d.largeSize {
		val := heap.Pop(&d.large).(int)

		d.largeSize--
		d.largeSum -= int64(val)

		heap.Push(&d.small, val)
		d.smallSize++
		d.smallSum += int64(val)

		d.pruneLarge()
	}
}

func (d *DualHeap) add(num int) {
	if d.small.Len() == 0 || num <= d.small[0] {
		heap.Push(&d.small, num)
		d.smallSize++
		d.smallSum += int64(num)
	} else {
		heap.Push(&d.large, num)
		d.largeSize++
		d.largeSum += int64(num)
	}

	d.rebalance()
}

func (d *DualHeap) remove(num int) {
	d.delayed[num]++

	if num <= d.small[0] {
		d.smallSize--
		d.smallSum -= int64(num)

		if num == d.small[0] {
			d.pruneSmall()
		}
	} else {
		d.largeSize--
		d.largeSum -= int64(num)

		if d.large.Len() > 0 && num == d.large[0] {
			d.pruneLarge()
		}
	}

	d.rebalance()
}

func (d *DualHeap) cost() int64 {
	median := int64(d.small[0])

	left := median*int64(d.smallSize) - d.smallSum
	right := d.largeSum - median*int64(d.largeSize)

	return left + right
}

func minOperations(nums []int, x int, k int) int64 {
	n := len(nums)

	dh := NewDualHeap()

	for i := 0; i < x; i++ {
		dh.add(nums[i])
	}

	windowCosts := make([]int64, 0, n-x+1)
	windowCosts = append(windowCosts, dh.cost())

	for r := x; r < n; r++ {
		dh.remove(nums[r-x])
		dh.add(nums[r])

		windowCosts = append(windowCosts, dh.cost())
	}

	const INF int64 = 1 << 60

	dp := make([][]int64, k+1)
	for i := range dp {
		dp[i] = make([]int64, n+1)
		for j := range dp[i] {
			dp[i][j] = INF
		}
	}

	dp[0][0] = 0

	for i := 0; i < n; i++ {
		for used := 0; used <= k; used++ {
			if dp[used][i] == INF {
				continue
			}

			if dp[used][i] < dp[used][i+1] {
				dp[used][i+1] = dp[used][i]
			}

			if used < k && i+x <= n {
				candidate := dp[used][i] + windowCosts[i]
				if candidate < dp[used+1][i+x] {
					dp[used+1][i+x] = candidate
				}
			}
		}
	}

	ans := INF
	for _, v := range dp[k] {
		if v < ans {
			ans = v
		}
	}

	return ans
}

Go-Specific Notes

The Go implementation mirrors the Python logic closely.

Because operation counts can become very large, all accumulated costs are stored as int64. Heap sums and DP values also use int64 to avoid overflow.

Go does not provide a built-in heap structure equivalent to Python's heapq, so the solution uses container/heap with custom min-heap and max-heap implementations.

Worked Examples

Example 1

nums = [5,-2,1,3,7,3,6,4,-1]
x = 3
k = 2

All length-3 windows:

Start Window
0 [5,-2,1]
1 [-2,1,3]
2 [1,3,7]
3 [3,7,3]
4 [7,3,6]
5 [3,6,4]
6 [6,4,-1]

Window costs:

Window Median Cost
[5,-2,1] 1 7
[-2,1,3] 1 5
[1,3,7] 3 6
[3,7,3] 3 4
[7,3,6] 6 4
[3,6,4] 4 3
[6,4,-1] 4 7

The DP considers all non-overlapping choices.

The optimal selection is:

start = 1, cost = 5
start = 5, cost = 3

Total:

5 + 3 = 8

Answer:

8

Example 2

nums = [9,-2,-2,-2,1,5]
x = 2
k = 2

Window costs:

Start Window Cost
0 [9,-2] 11
1 [-2,-2] 0
2 [-2,-2] 0
3 [-2,1] 3
4 [1,5] 4

DP chooses:

start = 1
start = 3

Total:

0 + 3 = 3

Answer:

3

Complexity Analysis

Measure Complexity Explanation
Time O(n log x + nk) Sliding-window median plus DP
Space O(n + kn) Window costs and DP table

The sliding-window median structure performs one insertion and one deletion per step, each costing O(log x). Since there are O(n) window movements, that phase costs O(n log x).

The dynamic programming phase has (k + 1) * (n + 1) states and constant-time transitions, giving O(nk) time.

Because k <= 15, the DP is easily fast enough even when n = 100000.

Test Cases

sol = Solution()

assert sol.minOperations(
    [5,-2,1,3,7,3,6,4,-1],
    3,
    2
) == 8  # example 1

assert sol.minOperations(
    [9,-2,-2,-2,1,5],
    2,
    2
) == 3  # example 2

assert sol.minOperations(
    [1,1,1,1],
    2,
    2
) == 0  # already valid

assert sol.minOperations(
    [1,2,3,4],
    2,
    1
) == 1  # single window

assert sol.minOperations(
    [10,10,10,10,10,10],
    3,
    2
) == 0  # all equal

assert sol.minOperations(
    [-5,-5,-5,-5],
    2,
    2
) == 0  # negative equal values

assert sol.minOperations(
    [1,100,1,100,1,100],
    2,
    3
) == 297  # every pair must be fixed

assert sol.minOperations(
    [1,2,1,2,1,2],
    2,
    2
) == 2  # choose two cheapest windows

assert sol.minOperations(
    [0,0,0,0,0,0],
    3,
    1
) == 0  # zero-cost window

assert sol.minOperations(
    [1,5,9,13,17,21],
    3,
    2
) == 16  # larger distances
Test Why
Example 1 Official sample
Example 2 Official sample
[1,1,1,1] Already satisfies requirement
[1,2,3,4] Smallest meaningful selection
All equal positive values Zero-cost scenario
All equal negative values Negative-number handling
Alternating large values Large modification costs
[1,2,1,2,1,2] Overlapping-window choices
All zeros Median cost should remain zero
Increasing sequence General nontrivial case

Edge Cases

Windows Already Equal

A window such as:

[7,7,7,7]

already satisfies the condition. The median is 7, and the sum of absolute deviations is 0. A buggy implementation might still perform unnecessary modifications or miscompute the median. The heap-based cost formula correctly returns zero because both left and right deviations are zero.

Large Negative and Positive Values

The input allows values as small as -10^6 and as large as 10^6. A window might contain both extremes. The implementation never relies on value ranges or coordinate compression. It only performs arithmetic on actual values and stores cumulative sums in 64-bit integers, preventing overflow.

Dense Overlapping Candidates

Consider:

nums = [1,1,1,1,1]
x = 2
k = 2

Every window has cost zero, but most windows overlap. A greedy strategy could accidentally count overlapping windows. The DP avoids this entirely because selecting a window at position i moves the state directly to i + x, ensuring future selections start outside the chosen interval.

Maximum Window Size

When:

x = n

there is only one possible window. The sliding-window logic still works because exactly one cost is generated, and the DP naturally handles the single available choice.

Multiple Optimal Solutions

Different sets of windows may produce the same minimum cost. The algorithm does not depend on reconstructing which windows were chosen. It only tracks the minimum achievable cost, so ties are handled automatically. The previous the official accepted solution, please provide the problem link or allow web verification first.