Divide an Array Into Subarrays With Minimum Cost II
Problem
Given nums of length n and positive integers k and dist, divide nums into k disjoint contiguous subarrays. The cost of a subarray is its first element. If the subarrays start at indices 0 = i₀ < i₁ < … < iₖ₋₁, the rule is iₖ₋₁ − i₁ ≤ dist. Return the minimum possible total cost nums[0] + nums[i₁] + … + nums[iₖ₋₁].
Since nums[0] is always included, we must pick k−1 more starts from [1, n−1] whose index spread is at most dist, minimizing the sum of nums at those indices.
nums = [1,3,2,6,4,2], k = 3, dist = 35def minimumCost(nums, k, dist):
n = len(nums)
m = k - 2 # how many "extra" starts live in the window
low, high = [], [] # low: max-heap of m smallest; high: the rest
where, removed = {}, set() # location flags + lazy-deletion set
low_sum = low_size = 0
def prune(h):
while h and h[0][1] in removed:
heapq.heappop(h)
def add(idx):
nonlocal low_sum, low_size
v = nums[idx]
prune(low)
if low_size < m or (low and v < -low[0][0]):
heapq.heappush(low, (-v, idx)); where[idx] = "low"
low_sum += v; low_size += 1
if low_size > m: # overflow: spill low's max to high
mv, mi = -low[0][0], low[0][1]
heapq.heappop(low)
heapq.heappush(high, (mv, mi)); where[mi] = "high"
low_sum -= mv; low_size -= 1
else:
heapq.heappush(high, (v, idx)); where[idx] = "high"
def remove(idx):
nonlocal low_sum, low_size
removed.add(idx)
if where.get(idx) == "low":
low_sum -= nums[idx]; low_size -= 1
def rebalance():
nonlocal low_sum, low_size
prune(low); prune(high)
while low_size < m and high: # refill low after deletions
v, idx = heapq.heappop(high)
heapq.heappush(low, (-v, idx)); where[idx] = "low"
low_sum += v; low_size += 1
prune(high)
while low and high and -low[0][0] > high[0][0]:
nv, ni = -low[0][0], low[0][1] # swap big low ↔ small high
hv, hi = heapq.heappop(high); heapq.heappop(low)
heapq.heappush(low, (-hv, hi)); where[hi] = "low"
heapq.heappush(high, (nv, ni)); where[ni] = "high"
low_sum += hv - nv; prune(low); prune(high)
best = float("inf"); hi = 0
for L in range(1, n):
Rt = min(L + dist, n - 1)
while hi < Rt: # grow window's right edge
hi += 1
if hi > L: add(hi)
remove(L) # L is the anchor; it leaves the pool
rebalance()
if low_size >= m:
best = min(best, nums[L] + low_sum)
return nums[0] + bestfunction minimumCost(nums, k, dist) {
const n = nums.length, m = k - 2; // m extra starts live in the window
const low = new MaxHeap(), high = new MinHeap();
const where = {}, removed = new Set();
let lowSum = 0, lowSize = 0;
function add(idx) {
const v = nums[idx];
low.prune(removed);
if (lowSize < m || (low.size() && v < low.top().val)) {
low.push({ val: v, idx }); where[idx] = "low";
lowSum += v; lowSize++;
if (lowSize > m) { // overflow: spill low's max to high
const t = low.pop(); high.push(t); where[t.idx] = "high";
lowSum -= t.val; lowSize--;
}
} else { high.push({ val: v, idx }); where[idx] = "high"; }
}
function remove(idx) {
removed.add(idx);
if (where[idx] === "low") { lowSum -= nums[idx]; lowSize--; }
}
function rebalance() {
low.prune(removed); high.prune(removed);
while (lowSize < m && high.size()) { // refill low after deletions
const t = high.pop(); low.push(t); where[t.idx] = "low";
lowSum += t.val; lowSize++; high.prune(removed);
}
while (low.size() && high.size() && low.top().val > high.top().val) {
const a = low.pop(), b = high.pop(); // swap big low with small high
low.push(b); where[b.idx] = "low";
high.push(a); where[a.idx] = "high";
lowSum += b.val - a.val;
low.prune(removed); high.prune(removed);
}
}
let best = Infinity, hi = 0;
for (let L = 1; L < n; L++) {
const Rt = Math.min(L + dist, n - 1);
while (hi < Rt) { hi++; if (hi > L) add(hi); } // grow right edge
remove(L); // L is the anchor; drop it
rebalance();
if (lowSize >= m) best = Math.min(best, nums[L] + lowSum);
}
return nums[0] + best;
}long minimumCost(int[] nums, int k, int dist) {
int n = nums.length, m = k - 2; // m extra starts live in the window
// low: max-heap of the m smallest; high: min-heap of the rest
PriorityQueue<int[]> low = new PriorityQueue<>((a,b)->b[0]-a[0]);
PriorityQueue<int[]> high = new PriorityQueue<>((a,b)->a[0]-b[0]);
Map<Integer,String> where = new HashMap<>();
Set<Integer> removed = new HashSet<>();
long lowSum = 0; int lowSize = 0;
long best = Long.MAX_VALUE; int hi = 0;
for (int L = 1; L < n; L++) {
int Rt = Math.min(L + dist, n - 1);
while (hi < Rt) { // grow window's right edge
hi++;
if (hi > L) {
int v = nums[hi];
while (!low.isEmpty() && removed.contains(low.peek()[1])) low.poll();
if (lowSize < m || (!low.isEmpty() && v < low.peek()[0])) {
low.add(new int[]{v, hi}); where.put(hi, "low");
lowSum += v; lowSize++;
if (lowSize > m) { // overflow: spill max to high
int[] t = low.poll(); high.add(t); where.put(t[1], "high");
lowSum -= t[0]; lowSize--;
}
} else { high.add(new int[]{v, hi}); where.put(hi, "high"); }
}
}
removed.add(L); // L is the anchor; drop it
if ("low".equals(where.get(L))) { lowSum -= nums[L]; lowSize--; }
while (!low.isEmpty() && removed.contains(low.peek()[1])) low.poll();
while (!high.isEmpty() && removed.contains(high.peek()[1])) high.poll();
while (lowSize < m && !high.isEmpty()) { // refill low after deletions
int[] t = high.poll();
if (removed.contains(t[1])) continue;
low.add(t); where.put(t[1], "low"); lowSum += t[0]; lowSize++;
}
while (!low.isEmpty() && !high.isEmpty() && low.peek()[0] > high.peek()[0]) {
int[] a = low.poll(), b = high.poll(); // swap big low with small high
low.add(b); where.put(b[1], "low");
high.add(a); where.put(a[1], "high");
lowSum += b[0] - a[0];
}
if (lowSize >= m) best = Math.min(best, nums[L] + lowSum);
}
return nums[0] + best;
}long long minimumCost(vector<int>& nums, int k, int dist) {
int n = nums.size(), m = k - 2; // m extra starts live in the window
// low: max-heap of m smallest; high: min-heap of the rest (value, index)
priority_queue<pair<int,int>> low;
priority_queue<pair<int,int>, vector<pair<int,int>>, greater<>> high;
unordered_map<int,int> where; // idx -> 1 (low) or 2 (high)
unordered_set<int> removed;
long long lowSum = 0; int lowSize = 0;
auto pruneLow = [&]{ while(!low.empty() && removed.count(low.top().second)) low.pop(); };
auto pruneHigh = [&]{ while(!high.empty() && removed.count(high.top().second)) high.pop(); };
long long best = LLONG_MAX; int hi = 0;
for (int L = 1; L < n; L++) {
int Rt = min(L + dist, n - 1);
while (hi < Rt) { // grow window's right edge
hi++;
if (hi > L) {
int v = nums[hi]; pruneLow();
if (lowSize < m || (!low.empty() && v < low.top().first)) {
low.push({v, hi}); where[hi] = 1; lowSum += v; lowSize++;
if (lowSize > m) { // overflow: spill max to high
auto t = low.top(); low.pop(); high.push(t); where[t.second] = 2;
lowSum -= t.first; lowSize--;
}
} else { high.push({v, hi}); where[hi] = 2; }
}
}
removed.insert(L); // L is the anchor; drop it
if (where.count(L) && where[L] == 1) { lowSum -= nums[L]; lowSize--; }
pruneLow(); pruneHigh();
while (lowSize < m && !high.empty()) { // refill low after deletions
auto t = high.top(); high.pop();
low.push(t); where[t.second] = 1; lowSum += t.first; lowSize++; pruneHigh();
}
while (!low.empty() && !high.empty() && low.top().first > high.top().first) {
auto a = low.top(); low.pop(); auto b = high.top(); high.pop(); // swap
low.push(b); where[b.second] = 1;
high.push(a); where[a.second] = 2;
lowSum += (long long)b.first - a.first; pruneLow(); pruneHigh();
}
if (lowSize >= m) best = min(best, (long long)nums[L] + lowSum);
}
return (long long)nums[0] + best;
}Explanation
The first subarray always starts at index 0, so nums[0] is a fixed part of the cost. The job is to pick k−1 more start indices from [1, n−1] whose index spread is at most dist, minimizing the sum of nums at those indices.
Fix the leftmost extra start L = i₁. The remaining k−2 starts must satisfy iₖ₋₁ ≤ L + dist, so they all lie in the window (L, min(L+dist, n−1)]. To minimize, we want the k−2 smallest values in that window; the answer for this L is nums[L] plus their sum.
As L slides right by one, the window's right edge grows and the index L itself leaves the candidate pool. We maintain the sum of the m = k−2 smallest values with two heaps: low, a max-heap holding exactly those m smallest (so low_sum is the answer contribution), and high, a min-heap holding everything else.
Inserting a value puts it into low when it belongs among the smallest, spilling low's maximum into high on overflow. Deletions are lazy: we mark an index removed and discard it only when it surfaces at a heap top. After each slide we rebalance() — refill low from high if deletions shrank it, then swap a too-large low top with a smaller high top until the partition is correct.
Example: nums = [1,3,2,6,4,2], k = 3, dist = 3. Here m = 1. The best leftmost start is L = 2 (value 2) with the single smallest of (2, 5] being index 5 (value 2), giving nums[0] + 2 + 2 = 5.