# Divide an Array Into Subarrays With Minimum Cost II
**Difficulty:** HARD
[External](https://leetcode.com/problems/divide-an-array-into-subarrays-with-minimum-cost-ii)
Canonical: https://scaleengineer.com/dsa/problems/divide-an-array-into-subarrays-with-minimum-cost-ii
**Patterns:** [Sliding Window](https://scaleengineer.com/dsa/patterns/sliding-window)
**Data structures:** Array, Hash Table, Heap (Priority Queue)
**Companies:** [American Express](https://scaleengineer.com/companies/american-express), [jio](https://scaleengineer.com/companies/jio)
---
## Problem
You are given a **0-indexed** array of integers `nums` of length `n`, and two **positive** integers `k` and `dist`.

The **cost** of an array is the value of its **first** element. For example, the cost of `[1,2,3]` is `1` while the cost of `[3,4,1]` is `3`.

You need to divide `nums` into `k` **disjoint contiguous** subarrays, such that the difference between the starting index of the **second** subarray and the starting index of the `kth` subarray should be **less than or equal to** `dist`. In other words, if you divide `nums` into the subarrays `nums[0..(i1 - 1)], nums[i1..(i2 - 1)], ..., nums[ik-1..(n - 1)]`, then `ik-1 - i1 <= dist`.

Return _the **minimum** possible sum of the cost of these_ _subarrays_.

**Example 1:**

**Input:** nums = [1,3,2,6,4,2], k = 3, dist = 3
**Output:** 5
**Explanation:** The best possible way to divide nums into 3 subarrays is: [1,3], [2,6,4], and [2]. This choice is valid because ik-1 - i1 is 5 - 2 = 3 which is equal to dist. The total cost is nums[0] + nums[2] + nums[5] which is 1 + 2 + 2 = 5.
It can be shown that there is no possible way to divide nums into 3 subarrays at a cost lower than 5.

**Example 2:**

**Input:** nums = [10,1,2,2,2,1], k = 4, dist = 3
**Output:** 15
**Explanation:** The best possible way to divide nums into 4 subarrays is: [10], [1], [2], and [2,2,1]. This choice is valid because ik-1 - i1 is 3 - 1 = 2 which is less than dist. The total cost is nums[0] + nums[1] + nums[2] + nums[3] which is 10 + 1 + 2 + 2 = 15.
The division [10], [1], [2,2,2], and [1] is not valid, because the difference between ik-1 and i1 is 5 - 1 = 4, which is greater than dist.
It can be shown that there is no possible way to divide nums into 4 subarrays at a cost lower than 15.

**Example 3:**

**Input:** nums = [10,8,18,9], k = 3, dist = 1
**Output:** 36
**Explanation:** The best possible way to divide nums into 4 subarrays is: [10], [8], and [18,9]. This choice is valid because ik-1 - i1 is 2 - 1 = 1 which is equal to dist.The total cost is nums[0] + nums[1] + nums[2] which is 10 + 8 + 18 = 36.
The division [10], [8,18], and [9] is not valid, because the difference between ik-1 and i1 is 3 - 1 = 2, which is greater than dist.
It can be shown that there is no possible way to divide nums into 3 subarrays at a cost lower than 36.

**Constraints:**

* `3 <= n <= 105`
* `1 <= nums[i] <= 109`
* `3 <= k <= n`
* `k - 2 <= dist <= n - 2`

# Approaches
## Iterating All Windows and Sorting
The core of the problem is to select `k-1` starting indices for the second to k-th subarrays from `nums[1:]`. The constraint `i_{k-1} - i_1 <= dist` implies that all these `k-1` chosen indices must lie within a contiguous block of `dist + 1` indices. To minimize the cost, we need to find a block (or window) of `dist + 1` indices in `[1, n-1]` such that the sum of the `k-1` smallest `nums` values within that block is minimized. This brute-force approach iterates through every possible window of size `dist + 1`, and for each window, it explicitly finds the `k-1` smallest elements by sorting, calculating their sum, and tracking the minimum sum found across all windows.
**Time:** O((n - dist) * dist * log(dist)). The outer loop runs about `n - dist` times. Inside the loop, we create a list of size `dist + 1` and sort it, which takes `O(dist * log(dist))` time. This is too slow for `n` and `dist` up to 10<sup>5</sup>. · **Space:** O(dist) to store the numbers in the current window for sorting.
**Pros:** Conceptually simple and easy to implement.; Correctly models the problem by checking all valid possibilities.
**Cons:** Highly inefficient due to repeated work.; Sorting within each sliding window step leads to a high time complexity, making it too slow for the given constraints.
### Explanation
The algorithm iterates through all possible starting positions for a window of size `dist + 1`. The windows slide over the part of the array from which we can choose our `k-1` start indices, which is `nums[1]` to `nums[n-1]`. 

Let `m = k - 1`. For each window, we do the following:
1. Extract the numbers within the current window into a temporary list.
2. Sort this list in non-decreasing order.
3. Sum the first `m` elements of the sorted list. This gives the minimum cost for the `k-1` subarrays if their start indices are chosen from this window.
4. We keep track of the minimum sum found among all windows.

Finally, the total minimum cost is `nums[0]` (the cost of the first subarray, which is fixed) plus the minimum sum we found. 

```java
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

class Solution {
    public long minimumCost(int[] nums, int k, int dist) {
        int n = nums.length;
        long minOtherCosts = Long.MAX_VALUE;
        int m = k - 1;

        // The window of candidate indices for the other k-1 subarrays
        // slides from [1, 1+dist] to [n-1-dist, n-1].
        // The starting index of the window 'i' can go from 1 to n-1-dist.
        for (int i = 1; i <= n - 1 - dist; i++) {
            List<Integer> windowNums = new ArrayList<>();
            for (int j = i; j <= i + dist; j++) {
                windowNums.add(nums[j]);
            }
            
            Collections.sort(windowNums);
            
            long currentSum = 0;
            for (int l = 0; l < m; l++) {
                currentSum += windowNums.get(l);
            }
            
            minOtherCosts = Math.min(minOtherCosts, currentSum);
        }
        
        return (long)nums[0] + minOtherCosts;
    }
}
```
### Algorithm
- The cost of the first subarray is always `nums[0]`, so we add it at the end. We need to find the minimum sum of costs for the other `k-1` subarrays.
- Let `m = k - 1`. We need to choose `m` indices from `nums[1:]`.
- The chosen indices must lie in a window of size `dist + 1`.
- Iterate through all possible start indices `i` for a window of size `dist + 1`. The window of indices is `[i, i + dist]`. `i` ranges from `1` to `n - 1 - dist`.
- For each window:
  - Collect all `nums[j]` for `j` from `i` to `i + dist` into a list.
  - Sort the list.
  - Calculate the sum of the first `m` elements.
  - Update a global minimum sum with this sum if it's smaller.
- Return `nums[0]` plus the global minimum sum.

## Optimized Sliding Window with Two Balanced BSTs (TreeMaps)
This approach significantly optimizes the previous method by avoiding redundant computations. Instead of re-sorting the entire window every time it slides, we use a more sophisticated data structure to maintain the `k-1` smallest elements efficiently. As the window slides one position to the right, one element is removed, and one is added. We can use two balanced binary search trees (implemented as `TreeMap` in Java to handle duplicates) to keep track of the elements in the window. One `TreeMap` (`smallMap`) will store the `k-1` smallest elements, and the other (`largeMap`) will store the rest. This allows us to find the sum of the `k-1` smallest elements in `O(log(dist))` time per slide, leading to a much faster overall solution.
**Time:** O(n log(dist)). The main loop iterates `O(n)` times. Each step involves an `add` and a `remove` operation on the `TreeMap`s. The maps store at most `dist + 1` distinct elements, so each map operation takes `O(log(dist))` time. · **Space:** O(dist) to store the elements of the window in the two `TreeMap`s.
**Pros:** Highly efficient with a time complexity suitable for the given constraints.; Effectively solves the problem by reducing re-computation in the sliding window.
**Cons:** Implementation is more complex and requires careful handling of the two-map data structure, including edge cases and balancing logic.; The logic for adding and removing elements while maintaining the two-map invariant can be tricky to get right.
### Explanation
We use a sliding window of size `dist + 1` over `nums[1:]`. To efficiently find the sum of the `k-1` smallest elements in this window, we use two `TreeMap`s: `smallMap` and `largeMap`.

- `smallMap`: Stores the `k-1` smallest elements from the current window. We also maintain `smallSum`, the sum of elements in `smallMap`, and `smallSize`, the count of elements.
- `largeMap`: Stores the remaining elements in the window.

The algorithm proceeds as follows:
1. Initialize the data structure with the elements from the first window, `nums[1...1+dist]`. We define `add` and `remove` helper methods that update the two maps and `smallSum` while keeping them balanced. An element is added or removed, and then elements are moved between the maps to ensure `smallMap` contains exactly the `k-1` smallest elements.
2. The initial `smallSum` is our first candidate for the minimum sum.
3. We then slide the window one element at a time from left to right. In each step, we `remove` the element that's leaving the window and `add` the element that's entering. 
4. After each slide, `smallSum` holds the sum of the `k-1` smallest elements in the new window. We update our overall minimum sum with the current `smallSum`.
5. The final answer is `nums[0]` plus this overall minimum sum.

This method reduces the complexity of each step of the sliding window from `O(dist * log(dist))` to `O(log(dist))`, making it efficient enough to pass the given constraints.

```java
import java.util.TreeMap;

class Solution {
    private TreeMap<Integer, Integer> smallMap;
    private TreeMap<Integer, Integer> largeMap;
    private long smallSum;
    private int smallSize;
    private int m; // Represents k - 1

    public long minimumCost(int[] nums, int k, int dist) {
        this.m = k - 1;
        this.smallMap = new TreeMap<>();
        this.largeMap = new TreeMap<>();
        this.smallSum = 0;
        this.smallSize = 0;

        for (int i = 1; i <= 1 + dist; i++) {
            add(nums[i]);
        }

        long minOtherCosts = smallSum;

        for (int i = 2; i <= nums.length - 1 - dist; i++) {
            remove(nums[i - 1]);
            add(nums[i + dist]);
            minOtherCosts = Math.min(minOtherCosts, smallSum);
        }

        return (long)nums[0] + minOtherCosts;
    }

    private void add(int num) {
        addToMap(smallMap, num);
        smallSum += num;
        smallSize++;
        balance();
    }

    private void remove(int num) {
        if (largeMap.containsKey(num)) {
            removeFromMap(largeMap, num);
        } else {
            removeFromMap(smallMap, num);
            smallSum -= num;
            smallSize--;
        }
        balance();
    }

    private void balance() {
        while (smallSize > m) {
            int maxInSmall = smallMap.lastKey();
            removeFromMap(smallMap, maxInSmall);
            smallSum -= maxInSmall;
            smallSize--;
            addToMap(largeMap, maxInSmall);
        }
        while (smallSize < m && !largeMap.isEmpty()) {
            int minInLarge = largeMap.firstKey();
            removeFromMap(largeMap, minInLarge);
            addToMap(smallMap, minInLarge);
            smallSum += minInLarge;
            smallSize++;
        }
    }

    private void addToMap(TreeMap<Integer, Integer> map, int num) {
        map.put(num, map.getOrDefault(num, 0) + 1);
    }

    private void removeFromMap(TreeMap<Integer, Integer> map, int num) {
        map.put(num, map.get(num) - 1);
        if (map.get(num) == 0) {
            map.remove(num);
        }
    }
}
```
### Algorithm
- Let `m = k - 1`. We need to find the minimum sum of `m` elements chosen from a sliding window of size `dist + 1` over `nums[1:]`.
- Use two `TreeMap`s: `smallMap` to store the `m` smallest elements and `largeMap` for the rest.
- Maintain `smallSum`, the sum of elements in `smallMap`.
- Initialize the window with elements from `nums[1]` to `nums[1 + dist]`. For each element, add it to the two-map structure and keep the maps balanced.
- After initialization, `smallSum` is the cost for the first window. Initialize `min_cost_sum` with this value.
- Iterate from `i = 2` to `n - 1 - dist` to slide the window:
  - Remove the element `nums[i-1]` from the data structure.
  - Add the new element `nums[i+dist]` to the data structure.
  - Both `add` and `remove` operations must maintain the properties of the two maps and update `smallSum`.
  - Update `min_cost_sum = min(min_cost_sum, smallSum)`.
- Return `nums[0] + min_cost_sum`.

# Solutions
### Java

```java
class Solution {
public
  long minimumCost(int[] nums, int k, int dist) {
    long result = Long.MAX_VALUE, sum = 0L;
    int n = nums.length;
    TreeSet<Integer> set1 =
        new TreeSet<>((a, b)->nums[a] == nums[b] ? a - b : nums[a] - nums[b]);
    TreeSet<Integer> set2 =
        new TreeSet<>((a, b)->nums[a] == nums[b] ? a - b : nums[a] - nums[b]);
    for (int i = 1; i < n; i++) {
      set1.add(i);
      sum += nums[i];
      if (set1.size() >= k) {
        int x = set1.pollLast();
        sum -= nums[x];
        set2.add(x);
      }
      if (i - dist > 0) {
        result = Math.min(result, sum);
        int temp = i - dist;
        if (set1.contains(temp)) {
          set1.remove(temp);
          sum -= nums[temp];
          if (set2.size() > 0) {
            int y = set2.pollFirst();
            sum += nums[y];
            set1.add(y);
          }
        } else {
          set2.remove(i - dist);
        }
      }
    }
    return result + nums[0];
  }
}

```

### CPP

```cpp
class Solution {
public:
  long long minimumCost(vector<int> &nums, int k, int dist) {
    multiset<int> sml, big;
    int sz = dist + 1;
    long long sum = 0, ans = 0;
    for (int i = 1; i <= sz; i++) {
      sml.insert(nums[i]);
      sum += nums[i];
    }
    while (sml.size() > k - 1) {
      big.insert(*sml.rbegin());
      sum -= *sml.rbegin();
      sml.erase(sml.find(*sml.rbegin()));
    }
    ans = sum;
    for (int i = sz + 1; i < nums.size(); i++) {
      sum += nums[i];
      sml.insert(nums[i]);
      if (big.find(nums[i - sz]) != big.end()) {
        big.erase(big.find(nums[i - sz]));
      } else {
        sum -= nums[i - sz];
        sml.erase(sml.find(nums[i - sz]));
      }
      while (sml.size() > k - 1) {
        sum -= *sml.rbegin();
        big.insert(*sml.rbegin());
        sml.erase(sml.find(*sml.rbegin()));
      }
      while (sml.size() < k - 1) {
        sum += *big.begin();
        sml.insert(*big.begin());
        big.erase(big.begin());
      }
      while (!sml.empty() && !big.empty() && *sml.rbegin() > *big.begin()) {
        sum -= *sml.rbegin() - *big.begin();
        sml.insert(*big.begin());
        big.insert(*sml.rbegin());
        sml.erase(sml.find(*sml.rbegin()));
        big.erase(big.begin());
      }
      ans = min(ans, sum);
    }
    int p = 0;
    return nums[0] + ans;
  }
};

```

### Python

```python
from sortedcontainers import SortedList class Solution : def minimumCost ( self , nums : List [ int ], k : int , dist : int ) -> int : n = len ( nums ) sl = SortedList () y = nums [ 0 ] ans = float ( "inf" ) i = 1 running_sum = 0 for j in range ( 1 , n ): pos = bisect . bisect_left ( sl , nums [ j ]) sl . add ( nums [ j ]) if pos < k - 1 : running_sum += nums [ j ] if len ( sl ) > k - 1 : running_sum -= sl [ k - 1 ] while j - i > dist : removed_pos = sl . index ( nums [ i ]) removed_element = nums [ i ] sl . remove ( removed_element ) if removed_pos < k - 1 : running_sum -= removed_element if len ( sl ) >= k - 1 : running_sum += sl [ k - 2 ] i += 1 if j - i + 1 >= k - 1 : ans = min ( ans , running_sum ) return ans + y
```