Divide an array into subarrays with minimum cost ii

Array Hash Table Sliding Window Heap (Priority Queue)

54.8% acceptance

Feb 25, 2026

543

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.

Solution

Rust

Time O(n log n)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn minimum_cost(nums: Vec<i32>, k: i32, dist: i32) -> i64 {
    use std::collections::BTreeMap;
    let n = nums.len();
    let k = k as usize;
    let dist = dist as usize;
    let need = k - 1;

    // Window [r-dist, r] over indices 1..n-1
    // Maintain k-1 smallest in window using two BTreeMaps
    let mut small: BTreeMap<i32, usize> = BTreeMap::new();
    let mut large: BTreeMap<i32, usize> = BTreeMap::new();
    let mut small_sum = 0i64;
    let mut small_count = 0usize;

    let mut ans = i64::MAX;

    for r in 1..n {
      // Add nums[r]
      Self::add(&mut small, &mut large, &mut small_sum, &mut small_count, need, nums[r]);
      // Remove element leaving window: nums[r - dist - 1] if r >= dist + 2
      if r >= dist + 2 {
        Self::rem(&mut small, &mut large, &mut small_sum, &mut small_count, nums[r - dist - 1]);
      }
      // Record if we have enough elements
      if small_count == need {
        ans = ans.min(nums[0] as i64 + small_sum);
      }
    }
    ans
  }

  fn add(small: &mut std::collections::BTreeMap<i32,usize>, large: &mut std::collections::BTreeMap<i32,usize>,
       small_sum: &mut i64, small_count: &mut usize, need: usize, v: i32) {
    if *small_count < need {
      *small.entry(v).or_insert(0) += 1;
      *small_sum += v as i64;
      *small_count += 1;
    } else if let Some((&max_s, _)) = small.range(..).next_back() {
      if v < max_s {
        *small.entry(v).or_insert(0) += 1;
        *small_sum += v as i64;
        let e = small.entry(max_s).or_insert(0);
        *e -= 1; if *e == 0 { small.remove(&max_s); }
        *small_sum -= max_s as i64;
        *large.entry(max_s).or_insert(0) += 1;
      } else {
        *large.entry(v).or_insert(0) += 1;
      }
    } else {
      *large.entry(v).or_insert(0) += 1;
    }
  }

  fn rem(small: &mut std::collections::BTreeMap<i32,usize>, large: &mut std::collections::BTreeMap<i32,usize>,
       small_sum: &mut i64, small_count: &mut usize, v: i32) {
    if small.get(&v).copied().unwrap_or(0) > 0 {
      let e = small.entry(v).or_insert(0);
      *e -= 1; if *e == 0 { small.remove(&v); }
      *small_sum -= v as i64;
      *small_count -= 1;
      if let Some((&min_l, _)) = large.iter().next() {
        let e = large.entry(min_l).or_insert(0);
        *e -= 1; if *e == 0 { large.remove(&min_l); }
        *small.entry(min_l).or_insert(0) += 1;
        *small_sum += min_l as i64;
        *small_count += 1;
      }
    } else {
      let e = large.entry(v).or_insert(0);
      *e -= 1; if *e == 0 { large.remove(&v); }
    }
  }
}