We use cookies (including Google cookies) to personalize ads and analyze traffic. By continuing to use our site, you accept our Privacy Policy.

SimplyLeet
TagsCompaniesDisclaimersPrivacyTermsContactAbout Us

Divide an Array Into Subarrays With Minimum Cost II

Number: 3260

Difficulty: Hard

Paid? No

Companies: Microsoft, American Express

Problem Description

Given a 0-indexed integer array nums along with positive integers k and dist, you must “cut” nums into k disjoint contiguous subarrays. The cost of a subarray is defined as the value of its first element. In any valid division, if we denote the starting indices of the k subarrays as 0, i₁, i₂, …, iₖ₋₁ then an extra constraint requires that the difference between the starting index of the second subarray and the starting index of the kᵗʰ subarray is at most dist (i.e. iₖ₋₁ – i₁ ≤ dist). Return the minimum possible total cost (that is, the sum of the first elements of the chosen subarrays).

For example, if nums = [1,3,2,6,4,2], k = 3 and dist = 3 then one optimal solution is to form subarrays [1,3], [2,6,4], and [2] with cost = 1 + 2 + 2 = 5, and note that the difference between the start indices of the second ([2,6,4] starting at index 2) and third ([2] starting at index 5) subarrays is 3.

Key Insights

  • We must choose k subarrays; the cost of a division is the sum of the first elements of these subarrays.
  • The cut-points (starting indices of subarrays, except the first which is fixed at index 0) can be interpreted as choosing an increasing sequence i₁, i₂, …, iₖ₋₁.
  • The additional constraint is iₖ₋₁ – i₁ ≤ dist.
  • An equivalent formulation is: fix the starting index of the second subarray as i₁; then choose k – 2 additional indices from (i₁+1 to min(n – 1, i₁+dist)] such that the overall cost (nums[0] + nums[i₁] + …) is minimized.
  • In absence of order–related “penalties” (beyond “first‐element” cost) the subarrays’ cost simply sum up the chosen starting elements.
  • A dynamic programming formulation can be defined where the state represents “how many subarray starting indices have been chosen so far” along with extra “memory” of the first cut (i₁) so that when picking the kᵗʰ subarray we can enforce the constraint.
  • To allow for efficient transitions (since n can be as large as 10⁵) one can use data structures such as segment trees (or heaps combined with sliding window techniques) to quickly query the minimum cost over a range.
  • One may “layer” DP for the subarrays starting from index 0, then for each candidate i₁ (for the second subarray) run a transition that “extends” the chain up to the kᵗʰ subarray while keeping track of the minimum cost and propagating along the “first‐cut” (i₁) value. Finally, the answer is the minimum over all DP states for which the chosen final index iₖ₋₁ satisfies iₖ₋₁ – i₁ ≤ dist.

Space and Time Complexity

Time Complexity: O(k · n · log n) in a solution that uses a segment tree (or alternatively similar time with a heap+sliding window) for optimizing the DP transitions. Space Complexity: O(k · n)

Note: With careful observation and further optimizations the time may be reduced; however, due to the extra constraint the DP state must propagate the first chosen subarray start and enforce it when selecting the final index.

Solution

The idea is to use dynamic programming where we define a state for the “chain” of subarray starting indices. In a straightforward (but too slow) formulation one might write: • Let dp[r][i] be the minimum cost to pick r subarrays ending with the rᵗʰ subarray starting at index i. • For the first subarray, dp[1][0] = nums[0]. • For r = 2, one can only “cut” at some index i > 0 and the cost becomes dp[2][i] = nums[0] + nums[i]. We then “store” i as the first cut value. • For r ≥ 3, we have dp[r][i] = nums[i] + min{ dp[r–1][j] } for j < i, and when propagating the state, we keep along the candidate first cut (which for r = 2 is simply j itself and for r > 2 is inherited). • In the end, we need to take the minimum dp[k][i] such that i – (first-cut value) ≤ dist. Because straightforward transitions will be too slow (O(n²) per layer), we use a segment tree (or similar range–minimum query DS) to quickly obtain the minimum over the range of valid indices from the previous DP layer. In each DP layer the transition is essentially “for each index i we query over all indices j < i (from the previous layer) to get min dp[r–1][j] and then add nums[i].” When r equals 2, the state “remembers” the candidate i (since dp[2][i] starts with first cut equal to i). For deeper layers (r ≥ 3), this “first cut” is propagated unchanged. Finally, the answer is computed as the minimum dp[k][i] with the additional validity check that i – (first-cut value stored in dp state) ≤ dist.

The code solutions below implement this idea with detailed inline comments. (In an interview you might explain that while the “naïve” DP is O(n²) time, using segment trees to maintain and query the minimum among DP states allows an efficient O(k · n · log n) solution.)

Code Solutions

# Python solution using DP with a segment tree.
# Note: This implementation provides the conceptual outline.
# In practice one must carefully implement an efficient segment tree.

import math

# Define a segment tree class for range minimum queries over DP states.
class SegmentTree:
    def __init__(self, arr):
        self.n = len(arr)
        self.size = 1 << (self.n-1).bit_length()
        self.tree = [math.inf]*(2*self.size)
        # We also keep an auxiliary array "first_cut"
        # that holds the first subarray starting index for each DP state.
        self.first_cut = [None]*(2*self.size)
        # Build the tree using the base array:
        for i in range(self.n):
            self.tree[self.size+i] = arr[i][0]  # dp cost is stored at index 0 of tuple (cost, first_cut)
            self.first_cut[self.size+i] = arr[i][1]
        for i in range(self.n, self.size):
            self.tree[self.size+i] = math.inf
            self.first_cut[self.size+i] = None
        for i in range(self.size-1, 0, -1):
            left = 2*i
            right = 2*i+1
            # For tie breaking, simply choose the one with lower cost.
            if self.tree[left] <= self.tree[right]:
                self.tree[i] = self.tree[left]
                self.first_cut[i] = self.first_cut[left]
            else:
                self.tree[i] = self.tree[right]
                self.first_cut[i] = self.first_cut[right]
                
    # Update position pos with new value (dp_cost, first_cut) pair.
    def update(self, pos, val):
        pos += self.size
        self.tree[pos] = val[0]
        self.first_cut[pos] = val[1]
        pos //= 2
        while pos:
            left = 2*pos
            right = left+1
            if self.tree[left] <= self.tree[right]:
                self.tree[pos] = self.tree[left]
                self.first_cut[pos] = self.first_cut[left]
            else:
                self.tree[pos] = self.tree[right]
                self.first_cut[pos] = self.first_cut[right]
            pos //= 2

    # Query the minimum dp state in [l, r).
    def query(self, l, r):
        min_val = math.inf
        min_first_cut = None
        l += self.size; r += self.size
        while l < r:
            if l & 1:
                if self.tree[l] < min_val:
                    min_val = self.tree[l]
                    min_first_cut = self.first_cut[l]
                l += 1
            if r & 1:
                r -= 1
                if self.tree[r] < min_val:
                    min_val = self.tree[r]
                    min_first_cut = self.first_cut[r]
            l //= 2; r //= 2
        return (min_val, min_first_cut)

def minCostDivideArray(nums, k, dist):
    n = len(nums)
    math_inf = math.inf
    # dp[r][i] will be stored only for indices i that are valid subarray starting points
    # We will use a list of pairs (cost, first_cut) for the current DP layer.
    # For r = 1: only index 0 is valid.
    dp_prev = [(math_inf, None)] * n
    dp_prev[0] = (nums[0], None)  # cost and no first cut for the first subarray

    # For r = 2: we choose one cut from index 1 to n-1.
    dp_curr = [(math_inf, None)] * n
    for i in range(1, n):
        # When starting the second subarray at index i, cost = dp_prev[0] + nums[i]
        dp_curr[i] = (dp_prev[0][0] + nums[i], i)
    dp_prev = dp_curr

    # For r from 3 to k, update DP.
    for r in range(3, k+1):
        dp_curr = [(math_inf, None)] * n
        # Build segment tree for dp_prev
        seg = SegmentTree(dp_prev)
        # We can only choose an index i where there is at least one index before.
        for i in range(r-1, n):
            # Query over indices [0, i) from previous layer
            (min_cost, first_cut) = seg.query(0, i)
            if min_cost < math_inf:
                # Cost to choose index i as new cut
                dp_curr[i] = (min_cost + nums[i], first_cut)
        dp_prev = dp_curr

    # Finally, among dp_prev for r == k, select the one with i - first_cut <= dist.
    ans = math_inf
    for i in range(k-1, n):
        cost, fcut = dp_prev[i]
        # fcut was set when r==2; ensure the constraint holds.
        if fcut is not None and (i - fcut) <= dist:
            ans = min(ans, cost)
    return ans if ans < math_inf else -1

# Example usage:
#print(minCostDivideArray([1,3,2,6,4,2], 3, 3))  # Expected output: 5
← Back to All Questions