3013. Divide an Array Into Subarrays With Minimum Cost II


Problem Description

You are tasked with taking an array of integers called nums, which has n elements, and you need to divide it into k disjoint contiguous subarrays. The unique challenge here is that the starting indices of the k subarrays must be such that the starting index of the second subarray and the kth subarray are no more than dist units apart. The cost of an array is determined by its first element, and your goal is to minimize the sum of the costs for these k subarrays.

To illustrate, consider you split nums into subarrays like so: nums[0..(i1 - 1)], nums[i1..(i2 - 1)], ..., nums[i(k-1)..(n - 1)], where each i represents the starting index of a subarray. The condition you must satisfy is that i(k-1) - i1 <= dist. The problem asks for the minimum possible sum of the first elements from these subarrays.

Intuition

The key to solving this problem is to focus on how we choose the starting points of the k subarrays to minimize the sum of their first elements while complying with the condition that the distance between the starting indices of the second and the kth subarray is no more than dist.

One strategy is to use a sliding window approach to address the constraints related to the indices and their distances. We keep track of a sorted list of values inside our current window, which can be done efficiently using data structures like a SortedList in Python.

At every step, we:

  1. Add the current element from the nums array to our sorted list.
  2. If our list's size is smaller than k-1, we add the current element to our running sum.
  3. If the list exceeds k-1, we remove the 'k-th' element from our running sum since it will not be the smallest when considering the first k-1 elements.
  4. We adjust the window if the distance between the start and the current element exceeds dist, removing elements from the beginning of the window and adjusting our running sum accordingly.

We are constantly looking for the smallest k-1 elements within our window at any given time, so when we reach the end of the array, we can be sure that we've explored all possible valid subarray combinations, and then we simply add the first element of nums to the final answer (since the cost of the first subarray is always going to be nums[0]).

This way, the running sum, at any given time, holds the sum of the smallest k-1 elements within a valid window, and we can update our answer by comparing it with the current running sum. Finally, we return the smallest sum found during the process added to the cost of the first element.

Learn more about Sliding Window and Heap (Priority Queue) patterns.

Not Sure What to Study? Take the 2-min Quiz to Find Your Missing Piece:

What's the output of running the following function using input [30, 20, 10, 100, 33, 12]?

1def fun(arr: List[int]) -> List[int]:
2    import heapq
3    heapq.heapify(arr)
4    res = []
5    for i in range(3):
6        res.append(heapq.heappop(arr))
7    return res
8
1public static int[] fun(int[] arr) {
2    int[] res = new int[3];
3    PriorityQueue<Integer> heap = new PriorityQueue<>();
4    for (int i = 0; i < arr.length; i++) {
5        heap.add(arr[i]);
6    }
7    for (int i = 0; i < 3; i++) {
8        res[i] = heap.poll();
9    }
10    return res;
11}
12
1class HeapItem {
2    constructor(item, priority = item) {
3        this.item = item;
4        this.priority = priority;
5    }
6}
7
8class MinHeap {
9    constructor() {
10        this.heap = [];
11    }
12
13    push(node) {
14        // insert the new node at the end of the heap array
15        this.heap.push(node);
16        // find the correct position for the new node
17        this.bubble_up();
18    }
19
20    bubble_up() {
21        let index = this.heap.length - 1;
22
23        while (index > 0) {
24            const element = this.heap[index];
25            const parentIndex = Math.floor((index - 1) / 2);
26            const parent = this.heap[parentIndex];
27
28            if (parent.priority <= element.priority) break;
29            // if the parent is bigger than the child then swap the parent and child
30            this.heap[index] = parent;
31            this.heap[parentIndex] = element;
32            index = parentIndex;
33        }
34    }
35
36    pop() {
37        const min = this.heap[0];
38        this.heap[0] = this.heap[this.size() - 1];
39        this.heap.pop();
40        this.bubble_down();
41        return min;
42    }
43
44    bubble_down() {
45        let index = 0;
46        let min = index;
47        const n = this.heap.length;
48
49        while (index < n) {
50            const left = 2 * index + 1;
51            const right = left + 1;
52
53            if (left < n && this.heap[left].priority < this.heap[min].priority) {
54                min = left;
55            }
56            if (right < n && this.heap[right].priority < this.heap[min].priority) {
57                min = right;
58            }
59            if (min === index) break;
60            [this.heap[min], this.heap[index]] = [this.heap[index], this.heap[min]];
61            index = min;
62        }
63    }
64
65    peek() {
66        return this.heap[0];
67    }
68
69    size() {
70        return this.heap.length;
71    }
72}
73
74function fun(arr) {
75    const heap = new MinHeap();
76    for (const x of arr) {
77        heap.push(new HeapItem(x));
78    }
79    const res = [];
80    for (let i = 0; i < 3; i++) {
81        res.push(heap.pop().item);
82    }
83    return res;
84}
85

Solution Approach

The solution employs a two-pointer strategy commonly used in sliding window problems, combined with a SortedList to maintain the contiguous elements in sorted order. This approach effectively balances between the constraints of the problem requirement. Here is a step by step explanation of how the algorithm is implemented:

  1. Initialize Pointers and Data Structures:

    • sl (SortedList): Helps to keep track of the sorted elements within the current window for efficient addition and removal of elements.
    • y (int): Stores the cost of the first subarray, which is always nums[0].
    • ans (int): Initializes with float("inf") to hold the minimum cost found.
    • i (int): Represents the starting index of the current window in the nums array.
    • running_sum (int): Accumulator to store the sum of the costs of the first k-1 smallest elements within the window.
  2. Iterate Through the Array:

    • Loop through the array starting from the second element because the first element is fixed as the first subarray's cost.
    • Get the position where the current element (nums[j]) should be inserted in sl using binary search.
    • Add the current element to sl.
  3. Maintain running_sum and sl:

    • When our current window has fewer than k-1 elements, add every newly encountered element to running_sum.
    • Once we have more than k-1 elements, remove the k-th smallest element from running_sum, as we are only interested in the sum of the first k-1 smallest elements.
  4. Adjust the Start Pointer:

    • The two-pointer approach is used to keep the dist constraint in check.
    • If the current window size exceeds dist, increment the starting index i and adjust the sl and running_sum accordingly.
    • The element at the start of the current window is removed from sl.
    • The running sum is updated by removing the removed element’s cost, and if necessary, adding the cost of the new k-2th smallest element to maintain the first k-1 smallest elements' sum.
  5. Update the Answer:

    • After adjusting the window, if the window has at least k-1 elements, we have a new potential solution.
    • Compare the current running_sum with ans and update ans if the current sum is smaller.
  6. Return the Result:

    • Since the first subarray cost is always nums[0], add it to the ans to get the minimum possible sum of the costs.

The code loops through each element in the array only once, making the runtime complexity O(n log k) given that insertion and removal operations in a SortedList have a logarithmic time complexity and we perform them at most n times. This approach strategically satisfies all of the problem's constraints, resulting in an efficient and effective solution.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Which technique can we use to find the middle of a linked list?

Example Walkthrough

Let's take an example to walk through the algorithm described above. Assume nums = [5, 2, 4, 1, 6, 3], k = 3, and dist = 3.

  1. Initialize Pointers and Data Structures:

    • sl (SortedList) will be used to maintain the sorted order of elements in the current window.
    • y is nums[0], which is 5 in our case.
    • ans is initialized to float("inf"), representing the minimum cost we are looking for.
    • i is set to 1, since we start from the second element (index 1).
    • running_sum starts at 0.
  2. Iterate Through the Array:

    • Start the loop at j = 1 with element 2 in nums.
  3. Maintain running_sum and sl:

    • Insert 2 into sl which is now [2]. running_sum becomes 2 (since we have fewer than k-1 elements).
    • Move to the next element, 4, insert it into sl which becomes [2, 4]. running_sum is 2 + 4 = 6.
    • Move to 1, insert into sl to have [1, 2, 4]. running_sum must now only include the smallest k-1 = 2 elements, which are [1, 2], so running_sum is updated to 1 + 2 = 3.
  4. Adjust the Start Pointer:

    • Now at index 3, the element 1 has been processed, and we're looking to process element 6, which is at index 4. We're within the dist constraint (index 4 - index 1 <= dist), so no adjustments are needed.
  5. Update the Answer:

    • We now have k-1 = 2 elements in our window and running_sum is 3. Update ans from inf to 3 + y = 3 + 5 = 8.

Proceeding with the same steps till the end:

  • When at element 3 (index 5), we would have a sl of [1, 2, 3, 4]. We ensure the sum of the k-1 smallest elements 1 + 2 = 3 is in running_sum. Since our window is too large (index 5 - index 1 = 4 > dist), we need to pop from sl and update the window start. Remove 2 (at window start) from sl becoming [1, 3, 4], adjust running_sum to include 1 + 3, and move i forward.

  • With the running sum now being 4 and we still satisfy the window size index 5 - index 2 = 3 <= dist, update ans to 4 + y = 4 + 5 = 9 if it's smaller than the previous ans.

  1. Return the Result:
    • After processing all the elements, we found the minimum possible sum of the costs to be 8 (from our prior step). This 8 includes the cost of the first element as per the problem statement.
    • The final answer is 8, representing the sum of the first elements of the k disjoint contiguous subarrays.

By following these steps, we can efficiently solve the problem using a sliding window approach combined with a sorted list to keep track of the elements and their costs within the window.

Solution Implementation

1from sortedcontainers import SortedList
2import bisect
3
4class Solution:
5    def minimumCost(self, nums: List[int], k: int, dist: int) -> int:
6        # Length of the input numbers list
7        num_length = len(nums)
8
9        # Initialize a sorted list to maintain a sorted order of numbers
10        sorted_list = SortedList()
11      
12        # The first element in the nums list as the starting point
13        starting_num = nums[0]
14      
15        # Initialize the answer to infinite for future comparison
16        min_cost = float("inf")
17      
18        # The window starts from the first element
19        left_index = 1
20      
21        # Running sum of the k-1 elements
22        running_sum = 0
23
24        # Loop through the nums list starting from the second element
25        for right_index in range(1, num_length):
26            # Find the position where nums[right_index] should be inserted
27            insert_pos = bisect.bisect_left(sorted_list, nums[right_index])
28          
29            # Add the current number to the sorted list
30            sorted_list.add(nums[right_index])
31
32            # If the insert position is less than k - 1, update the running sum
33            if insert_pos < k - 1:
34                running_sum += nums[right_index]
35                # If the sorted list has more than k - 1 elements, subtract the k-th element
36                if len(sorted_list) > k - 1:
37                    running_sum -= sorted_list[k - 1]
38
39            # If the window size is larger than allowed by 'dist', adjust the window from the left
40            while right_index - left_index > dist:
41                # Find the index of the left-most number to be removed
42                removed_index = sorted_list.index(nums[left_index])
43                # Remove the left-most element
44                removed_num = nums[left_index]
45                sorted_list.remove(removed_num)
46
47                # Adjust the running sum based on the position of the removed element
48                if removed_index < k - 1:
49                    running_sum -= removed_num
50                    # If there are still at least k - 1 elements, add the (k-1)-th element to the running sum
51                    if len(sorted_list) >= k - 1:
52                        running_sum += sorted_list[k - 2]
53                # Move the left window index to the right
54                left_index += 1
55
56            # If the window size is at least k - 1, calculate the cost to consider this subarray
57            if right_index - left_index + 1 >= k - 1:
58                min_cost = min(min_cost, running_sum)
59
60        # Return the minimum cost added with the first element that was set aside
61        return min_cost + starting_num
62
1class Solution {
2    public long minimumCost(int[] nums, int k, int dist) {
3        long minCost = Long.MAX_VALUE; // Initialize the minimum cost to the maximum possible value.
4        long currentSum = 0L;
5        int arrayLength = nums.length;
6      
7        // Create two TreeSets to manage the window of numbers within the distance 'dist'.
8        // The first TreeSet (setLower) keeps the smallest 'k' numbers seen so far.
9        // The second TreeSet (setHigher) keeps the rest.
10        // Custom comparator is used to sort the elements based on their values in 'nums'.
11        // In case of a tie, indices are compared to guarantee uniqueness.
12        TreeSet<Integer> setLower = new TreeSet<>((a, b) -> nums[a] == nums[b] ? a - b : nums[a] - nums[b]);
13        TreeSet<Integer> setHigher = new TreeSet<>((a, b) -> nums[a] == nums[b] ? a - b : nums[a] - nums[b]);
14      
15        // Start from the first index since the base case will always include nums[0].
16        for (int i = 1; i < arrayLength; i++) {
17            setLower.add(i); // Add the current element to the set of smaller elements.
18            currentSum += nums[i];
19          
20            // If there are more than 'k' elements in the lower set, move the largest one to the higher set.
21            if (setLower.size() > k) {
22                int highestInLower = setLower.pollLast();
23                currentSum -= nums[highestInLower];
24                setHigher.add(highestInLower);
25            }
26          
27            // Check if we should compute the cost for the current window.
28            if (i - dist > 0) {
29                minCost = Math.min(minCost, currentSum); // Update the minimum cost if needed.
30                int indexToRemove = i - dist; // Compute the index to remove from the current window.
31              
32                // Remove the element at 'indexToRemove' from the appropriate set.
33                if (setLower.contains(indexToRemove)) {
34                    setLower.remove(indexToRemove);
35                    currentSum -= nums[indexToRemove];
36                  
37                    // If higher set has elements, move the smallest one to the lower set to maintain 'k' elements.
38                    if (!setHigher.isEmpty()) {
39                        int lowestInHigher = setHigher.pollFirst();
40                        currentSum += nums[lowestInHigher];
41                        setLower.add(lowestInHigher);
42                    }
43                } else {
44                    setHigher.remove(indexToRemove);
45                }
46            }
47        }
48      
49        // Return the minimum cost found plus the base cost (nums[0]).
50        return minCost + nums[0];
51    }
52}
53
1class Solution {
2public:
3    // This function calculates the minimum sum of 'k' elements from each consecutive window of 'dist + 1' elements in the 'nums' vector.
4    long long minimumCost(vector<int>& nums, int k, int dist) {
5        // The small_set keeps the smallest 'k-1' elements, and the large_set keeps the others.
6        multiset<int> small_set, large_set;
7        int window_size = dist + 1;
8        long long window_sum = 0;  // Sum of the smallest 'k' elements in the current window.
9        long long min_sum = 0;     // Minimum sum found of 'k' elements across all windows.
10
11        // Initialize the first window
12        for (int i = 1; i <= window_size; i++) {
13            small_set.insert(nums[i]);
14            window_sum += nums[i];
15        }
16
17        // Ensure small_set only has 'k-1' smallest values
18        while (small_set.size() > k - 1) {
19            large_set.insert(*small_set.rbegin());
20            window_sum -= *small_set.rbegin();
21            small_set.erase(small_set.find(*small_set.rbegin()));
22        }
23        min_sum = window_sum;
24
25        // Slide the window across the array
26        for (int i = window_size + 1; i < nums.size(); i++) {
27            window_sum += nums[i];
28            small_set.insert(nums[i]);
29          
30            // Maintain the two multisets by moving the distanced element to the correct set
31            if (large_set.find(nums[i - window_size]) != large_set.end()) {
32                large_set.erase(large_set.find(nums[i - window_size]));
33            } else {
34                window_sum -= nums[i - window_size];
35                small_set.erase(small_set.find(nums[i - window_size]));
36            }
37
38            // Maintain the size of the small_set ('k-1' elements)
39            while (small_set.size() > k - 1) {
40                window_sum -= *small_set.rbegin();
41                large_set.insert(*small_set.rbegin());
42                small_set.erase(small_set.find(*small_set.rbegin()));
43            }
44
45            // Populate small_set if it has fewer than 'k-1' elements by taking from large_set
46            while (small_set.size() < k - 1) {
47                window_sum += *large_set.begin();
48                small_set.insert(*large_set.begin());
49                large_set.erase(large_set.begin());
50            }
51
52            // Balance the sets, so the elements in small_set are actually smaller than those in large_set
53            while (!small_set.empty() && !large_set.empty() && *small_set.rbegin() > *large_set.begin()) {
54                window_sum -= *small_set.rbegin() - *large_set.begin();
55                small_set.insert(*large_set.begin());
56                large_set.insert(*small_set.rbegin());
57                small_set.erase(small_set.find(*small_set.rbegin()));
58                large_set.erase(large_set.begin());
59            }
60          
61            min_sum = std::min(min_sum, window_sum); // Update the minimum sum found
62        }
63
64        // The result is the minimum sum of 'k' elements from each window plus the first element
65        return nums[0] + min_sum;
66    }
67};
68
1// Import the necessary collection class
2import { MultiSet } from "typescript-collections";
3
4// This function calculates the minimum sum of 'k' elements from each consecutive window of 'dist + 1' elements in the 'nums' array.
5function minimumCost(nums: number[], k: number, dist: number): bigint {
6    // The smallSet keeps the smallest 'k-1' elements, and the largeSet keeps the others.
7    const smallSet = new MultiSet<number>();
8    const largeSet = new MultiSet<number>();
9    const windowSize = dist + 1;
10    let windowSum: bigint = BigInt(0);  // Sum of the smallest 'k' elements in the current window.
11    let minSum: bigint = BigInt(0);     // Minimum sum found of 'k' elements across all windows.
12
13    // Initialize the first window
14    for (let i = 1; i <= windowSize; i++) {
15        smallSet.add(nums[i]);
16        windowSum += BigInt(nums[i]);
17    }
18
19    // Ensure smallSet only has 'k-1' smallest values
20    while (smallSet.size() > k - 1) {
21        largeSet.add(smallSet.higherElementsFirst()[0]);
22        windowSum -= BigInt(smallSet.higherElementsFirst()[0]);
23        smallSet.remove(smallSet.higherElementsFirst()[0]);
24    }
25    minSum = windowSum;
26
27    // Slide the window across the array
28    for (let i = windowSize + 1; i < nums.length; i++) {
29        windowSum += BigInt(nums[i]);
30        smallSet.add(nums[i]);
31      
32        // Maintain the two multisets by moving the distanced element to the correct set
33        if (largeSet.contains(nums[i - windowSize])) {
34            largeSet.remove(nums[i - windowSize]);
35        } else {
36            windowSum -= BigInt(nums[i - windowSize]);
37            smallSet.remove(nums[i - windowSize]);
38        }
39
40        // Maintain the size of the smallSet ('k-1' elements)
41        while (smallSet.size() > k - 1) {
42            windowSum -= BigInt(smallSet.higherElementsFirst()[0]);
43            largeSet.add(smallSet.higherElementsFirst()[0]);
44            smallSet.remove(smallSet.higherElementsFirst()[0]);
45        }
46
47        // Populate smallSet if it has fewer than 'k-1' elements by taking from largeSet
48        while (smallSet.size() < k - 1) {
49            windowSum += BigInt(largeSet.lowerElementsFirst()[0]);
50            smallSet.add(largeSet.lowerElementsFirst()[0]);
51            largeSet.remove(largeSet.lowerElementsFirst()[0]);
52        }
53
54        // Balance the sets so the elements in smallSet are actually smaller than those in largeSet
55        while (smallSet.size() > 0 && largeSet.size() > 0 && smallSet.higherElementsFirst()[0] > largeSet.lowerElementsFirst()[0]) {
56            windowSum += BigInt(largeSet.lowerElementsFirst()[0]);
57            windowSum -= BigInt(smallSet.higherElementsFirst()[0]);
58            smallSet.add(largeSet.lowerElementsFirst()[0]);
59            largeSet.add(smallSet.higherElementsFirst()[0]);
60            smallSet.remove(smallSet.higherElementsFirst()[0]);
61            largeSet.remove(largeSet.lowerElementsFirst()[0]);
62        }
63      
64        minSum = windowSum < minSum ? windowSum : minSum; // Update the minimum sum found
65    }
66
67    // The result is the minimum sum of 'k' elements from each window plus the first element
68    return BigInt(nums[0]) + minSum;
69}
70
71// You should also include the definitions for the MultiSet class or use an existing library
72// that implements this functionality.
73
Not Sure What to Study? Take the 2-min Quiz:

What is the best way of checking if an element exists in a sorted array once in terms of time complexity? Select the best that applies.

Time and Space Complexity

Time Complexity

The processed time complexity of the code is primarily determined by the main loop that runs through all elements in the list nums. Inside this loop, several operations are performed:

  1. Searching for the appropriate insertion position for the element in the sorted list: This is done with bisect.bisect_left(sl, nums[j]), which runs in O(log K) time, where K is the size of the sorted list, capped at k-1.

  2. Adding an element to the sorted list: The sl.add(nums[j]) operation also runs in O(log K) because it maintains the sorted property.

  3. Updating the running_sum involves simple arithmetic operations, which is O(1).

  4. Removing elements and updating the running_sum when the subarray exceeds the specified dist. Removal of an element is done with sl.remove(removed_element), preceded by finding the index of the element to be removed using sl.index(nums[i]). Indexing runs in O(K) time and removal runs in O(log K).

Given the loop runs n times and K can be at most k-1, the total time complexity is O(n * (2*log(K) + K)). Since K is bounded by k-1, we can simplify the expression to O(n * (log(k) + k)).

Space Complexity

The space complexity is affected by:

  1. The sorted list sl, which can contain at most k-1+dist elements (worst case scenario if all conditions never trigger a removal of elements before the dist constraint applies).

  2. Internal variables used for iteration and computation.

The dominant factor here is the sorted list. Given it can hold up to k-1+dist elements, the space complexity is O(K + dist). However, since K is limited to k-1, the space complexity simplifies to O(k + dist).

Learn more about how to find time and space complexity quickly using problem constraints.

Fast Track Your Learning with Our Quick Skills Quiz:

Which data structure is used to implement priority queue?


Recommended Readings


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫