Divide an Array Into Subarrays With Minimum Cost II

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 theminimum possible sum of the cost of these subarrays.

Examples

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 <= 10^5
• 1 <= nums[i] <= 10^9
• 3 <= k <= n
• k - 2 <= dist <= n - 2

Solution

Method 1 – Dynamic Programming with Sliding Window Minimum

Intuition

We want to split the array into k contiguous subarrays, minimizing the sum of their first elements, with the constraint that the difference between the starting index of the second and kth subarray is at most dist. We use DP where dp[i][j] is the minimum cost to split the first i elements into j subarrays, and use a sliding window minimum to efficiently find the best previous split for the constraint.

Approach

1. Let n be the length of nums. Initialize dp as a 2D array of size (n+1) x (k+1) with infinity, except dp[0][0] = 0.
2. For each i from 1 to n, and for each j from 1 to k:
   • For each possible previous split point p in [max(0, i-dist), i-1], update dp[i][j] = min(dp[i][j], dp[p][j-1] + nums[p]).
3. The answer is the minimum dp[i][k] for all valid i where the last subarray starts at i and the difference between the start of the second and kth subarray is at most dist.
4. To optimize, use a monotonic queue to maintain the minimum in the sliding window for each j.

Code

C++

class Solution {
public:
    int minimumCost(vector<int>& nums, int k, int dist) {
        int n = nums.size();
        vector<vector<long long>> dp(n + 1, vector<long long>(k + 1, 1e18));
        dp[0][0] = 0;
        for (int j = 1; j <= k; ++j) {
            deque<pair<long long, int>> dq;
            for (int i = 1; i <= n; ++i) {
                int l = max(0, i - dist);
                int r = i - 1;
                while (!dq.empty() && dq.front().second < l) dq.pop_front();
                if (r >= 0) {
                    long long val = dp[r][j - 1] + nums[r];
                    while (!dq.empty() && dq.back().first >= val) dq.pop_back();
                    dq.emplace_back(val, r);
                }
                if (!dq.empty()) dp[i][j] = dq.front().first;
            }
        }
        long long ans = 1e18;
        for (int i = 1; i <= n; ++i) ans = min(ans, dp[i][k]);
        return ans;
    }
};

Go

func minimumCost(nums []int, k int, dist int) int {
    n := len(nums)
    const inf = 1 << 60
    dp := make([][]int, n+1)
    for i := range dp {
        dp[i] = make([]int, k+1)
        for j := range dp[i] {
            dp[i][j] = inf
        }
    }
    dp[0][0] = 0
    for j := 1; j <= k; j++ {
        dq := [][2]int{}
        for i := 1; i <= n; i++ {
            l := i - dist
            if l < 0 {
                l = 0
            }
            r := i - 1
            for len(dq) > 0 && dq[0][1] < l {
                dq = dq[1:]
            }
            if r >= 0 {
                val := dp[r][j-1] + nums[r]
                for len(dq) > 0 && dq[len(dq)-1][0] >= val {
                    dq = dq[:len(dq)-1]
                }
                dq = append(dq, [2]int{val, r})
            }
            if len(dq) > 0 {
                dp[i][j] = dq[0][0]
            }
        }
    }
    ans := inf
    for i := 1; i <= n; i++ {
        if dp[i][k] < ans {
            ans = dp[i][k]
        }
    }
    return ans
}

Java

class Solution {
    public int minimumCost(int[] nums, int k, int dist) {
        int n = nums.length;
        long[][] dp = new long[n + 1][k + 1];
        for (long[] row : dp) Arrays.fill(row, (long)1e18);
        dp[0][0] = 0;
        for (int j = 1; j <= k; ++j) {
            Deque<long[]> dq = new ArrayDeque<>();
            for (int i = 1; i <= n; ++i) {
                int l = Math.max(0, i - dist);
                int r = i - 1;
                while (!dq.isEmpty() && dq.peekFirst()[1] < l) dq.pollFirst();
                if (r >= 0) {
                    long val = dp[r][j - 1] + nums[r];
                    while (!dq.isEmpty() && dq.peekLast()[0] >= val) dq.pollLast();
                    dq.offerLast(new long[]{val, r});
                }
                if (!dq.isEmpty()) dp[i][j] = dq.peekFirst()[0];
            }
        }
        long ans = (long)1e18;
        for (int i = 1; i <= n; ++i) ans = Math.min(ans, dp[i][k]);
        return (int)ans;
    }
}

Kotlin

class Solution {
    fun minimumCost(nums: IntArray, k: Int, dist: Int): Int {
        val n = nums.size
        val inf = 1e18.toLong()
        val dp = Array(n + 1) { LongArray(k + 1) { inf } }
        dp[0][0] = 0L
        for (j in 1..k) {
            val dq = ArrayDeque<Pair<Long, Int>>()
            for (i in 1..n) {
                val l = maxOf(0, i - dist)
                val r = i - 1
                while (dq.isNotEmpty() && dq.first().second < l) dq.removeFirst()
                if (r >= 0) {
                    val v = dp[r][j - 1] + nums[r]
                    while (dq.isNotEmpty() && dq.last().first >= v) dq.removeLast()
                    dq.addLast(v to r)
                }
                if (dq.isNotEmpty()) dp[i][j] = dq.first().first
            }
        }
        var ans = inf
        for (i in 1..n) ans = minOf(ans, dp[i][k])
        return ans.toInt()
    }
}

Python

class Solution:
    def minimumCost(self, nums: list[int], k: int, dist: int) -> int:
        n = len(nums)
        inf = float('inf')
        dp = [[inf] * (k + 1) for _ in range(n + 1)]
        dp[0][0] = 0
        for j in range(1, k + 1):
            from collections import deque
            dq: deque[tuple[int, int]] = deque()
            for i in range(1, n + 1):
                l = max(0, i - dist)
                r = i - 1
                while dq and dq[0][1] < l:
                    dq.popleft()
                if r >= 0:
                    val = dp[r][j - 1] + nums[r]
                    while dq and dq[-1][0] >= val:
                        dq.pop()
                    dq.append((val, r))
                if dq:
                    dp[i][j] = dq[0][0]
        return min(dp[i][k] for i in range(1, n + 1))

Rust

impl Solution {
    pub fn minimum_cost(nums: Vec<i32>, k: i32, dist: i32) -> i32 {
        let n = nums.len();
        let k = k as usize;
        let dist = dist as usize;
        let inf = i64::MAX / 2;
        let mut dp = vec![vec![inf; k + 1]; n + 1];
        dp[0][0] = 0;
        for j in 1..=k {
            let mut dq: std::collections::VecDeque<(i64, usize)> = std::collections::VecDeque::new();
            for i in 1..=n {
                let l = if i > dist { i - dist } else { 0 };
                let r = i - 1;
                while let Some(&(_, idx)) = dq.front() {
                    if idx < l { dq.pop_front(); } else { break; }
                }
                if r >= 0 {
                    let val = dp[r][j - 1] + nums[r - 0] as i64;
                    while let Some(&(v, _)) = dq.back() {
                        if v >= val { dq.pop_back(); } else { break; }
                    }
                    dq.push_back((val, r));
                }
                if let Some(&(v, _)) = dq.front() {
                    dp[i][j] = v;
                }
            }
        }
        let mut ans = inf;
        for i in 1..=n { ans = ans.min(dp[i][k]); }
        ans as i32
    }
}

TypeScript

class Solution {
    minimumCost(nums: number[], k: number, dist: number): number {
        const n = nums.length;
        const inf = 1e18;
        const dp: number[][] = Array.from({ length: n + 1 }, () => Array(k + 1).fill(inf));
        dp[0][0] = 0;
        for (let j = 1; j <= k; ++j) {
            const dq: [number, number][] = [];
            for (let i = 1; i <= n; ++i) {
                const l = Math.max(0, i - dist);
                const r = i - 1;
                while (dq.length && dq[0][1] < l) dq.shift();
                if (r >= 0) {
                    const val = dp[r][j - 1] + nums[r];
                    while (dq.length && dq[dq.length - 1][0] >= val) dq.pop();
                    dq.push([val, r]);
                }
                if (dq.length) dp[i][j] = dq[0][0];
            }
        }
        let ans = inf;
        for (let i = 1; i <= n; ++i) ans = Math.min(ans, dp[i][k]);
        return ans;
    }
}

Complexity

• ⏰ Time complexity: O(nk), each DP state is computed in amortized O(1) using a monotonic queue.
• 🧺 Space complexity: O(nk), for the DP table.