Question

3013. 将数组分成最小总代价的子数组 II

English Version

题目描述

给你一个下标从 0 开始长度为 n 的整数数组 nums 和两个 正 整数 k 和 dist 。

一个数组的 代价 是数组中的 第一个 元素。比方说，[1,2,3] 的代价为 1 ，[3,4,1] 的代价为 3 。

你需要将 nums 分割成 k 个 连续且互不相交 的子数组，满足 第二 个子数组与第 k 个子数组中第一个元素的下标距离 不超过 dist 。换句话说，如果你将 nums 分割成子数组 nums[0..(i1 - 1)], nums[i1..(i2 - 1)], ..., nums[ik-1..(n - 1)] ，那么它需要满足 ik-1 - i1 <= dist 。

请你返回这些子数组的 最小 总代价。

 

示例 1：

输入：nums = [1,3,2,6,4,2], k = 3, dist = 3
输出：5
解释：将数组分割成 3 个子数组的最优方案是：[1,3] ，[2,6,4] 和 [2] 。这是一个合法分割，因为 ik-1 - i1 等于 5 - 2 = 3 ，等于 dist 。总代价为 nums[0] + nums[2] + nums[5] ，也就是 1 + 2 + 2 = 5 。
5 是分割成 3 个子数组的最小总代价。

示例 2：

输入：nums = [10,1,2,2,2,1], k = 4, dist = 3
输出：15
解释：将数组分割成 4 个子数组的最优方案是：[10] ，[1] ，[2] 和 [2,2,1] 。这是一个合法分割，因为 ik-1 - i1 等于 3 - 1 = 2 ，小于 dist 。总代价为 nums[0] + nums[1] + nums[2] + nums[3] ，也就是 10 + 1 + 2 + 2 = 15 。
分割 [10] ，[1] ，[2,2,2] 和 [1] 不是一个合法分割，因为 ik-1 和 i1 的差为 5 - 1 = 4 ，大于 dist 。
15 是分割成 4 个子数组的最小总代价。

示例 3：

输入：nums = [10,8,18,9], k = 3, dist = 1
输出：36
解释：将数组分割成 4 个子数组的最优方案是：[10] ，[8] 和 [18,9] 。这是一个合法分割，因为 ik-1 - i1 等于 2 - 1 = 1 ，等于 dist 。总代价为 nums[0] + nums[1] + nums[2] ，也就是 10 + 8 + 18 = 36 。
分割 [10] ，[8,18] 和 [9] 不是一个合法分割，因为 ik-1 和 i1 的差为 3 - 1 = 2 ，大于 dist 。
36 是分割成 3 个子数组的最小总代价。

 

提示：

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

解法

方法一

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

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];
  }
}

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;
  }
};

func minimumCost(nums []int, k int, dist int) int64 {
  res := nums[0] + slices.Min(windowTopKSum(nums[1:], dist+1, k-1, true))
  return int64(res)
}

func windowTopKSum(nums []int, windowSize, k int, min bool) []int {
  n := len(nums)
  ts := NewTopKSum(k, min)
  res := []int{}
  for right := 0; right < n; right++ {
    ts.Add(nums[right])
    if right >= windowSize {
      ts.Discard(nums[right-windowSize])
    }
    if right >= windowSize-1 {
      res = append(res, ts.Query())
    }
  }
  return res
}

type TopKSum struct {
  sum   int
  k     int
  in    *Heap
  out   *Heap
  dIn   *Heap
  dOut  *Heap
  counter map[int]int
}

func NewTopKSum(k int, min bool) *TopKSum {
  var less func(a, b int) bool
  if min {
    less = func(a, b int) bool { return a < b }
  } else {
    less = func(a, b int) bool { return a > b }
  }
  return &TopKSum{
    k:     k,
    in:    NewHeap(less),
    out:   NewHeap(less),
    dIn:   NewHeap(less),
    dOut:  NewHeap(less),
    counter: map[int]int{},
  }
}

func (t *TopKSum) Query() int {
  return t.sum
}

func (t *TopKSum) Add(x int) {
  t.counter[x]++
  t.in.Push(-x)
  t.sum += x
  t.modify()
}

func (t *TopKSum) Discard(x int) bool {
  if t.counter[x] == 0 {
    return false
  }
  t.counter[x]--
  if t.in.Len() > 0 && -t.in.Top() == x {
    t.sum -= x
    t.in.Pop()
  } else if t.in.Len() > 0 && -t.in.Top() > x {
    t.sum -= x
    t.dIn.Push(-x)
  } else {
    t.dOut.Push(x)
  }
  t.modify()
  return true
}

func (t *TopKSum) SetK(k int) {
  t.k = k
  t.modify()
}

func (t *TopKSum) GetK() int {
  return t.k
}

func (t *TopKSum) Len() int {
  return t.in.Len() + t.out.Len() - t.dIn.Len() - t.dOut.Len()
}

func (t *TopKSum) Has(x int) bool {
  return t.counter[x] > 0
}

func (t *TopKSum) modify() {
  for t.out.Len() > 0 && (t.in.Len()-t.dIn.Len() < t.k) {
    p := t.out.Pop()
    if t.dOut.Len() > 0 && p == t.dOut.Top() {
      t.dOut.Pop()
    } else {
      t.sum += p
      t.in.Push(-p)
    }
  }

  for t.in.Len()-t.dIn.Len() > t.k {
    p := -t.in.Pop()
    if t.dIn.Len() > 0 && p == -t.dIn.Top() {
      t.dIn.Pop()
    } else {
      t.sum -= p
      t.out.Push(p)
    }
  }

  for t.dIn.Len() > 0 && t.in.Top() == t.dIn.Top() {
    t.in.Pop()
    t.dIn.Pop()
  }
}

type H = int

func NewHeap(less func(a, b H) bool, nums ...H) *Heap {
  nums = append(nums[:0:0], nums...)
  heap := &Heap{less: less, data: nums}
  heap.heapify()
  return heap
}

type Heap struct {
  data []H
  less func(a, b H) bool
}

func (h *Heap) Push(value H) {
  h.data = append(h.data, value)
  h.pushUp(h.Len() - 1)
}

func (h *Heap) Pop() (value H) {
  if h.Len() == 0 {
    panic("heap is empty")
  }

  value = h.data[0]
  h.data[0] = h.data[h.Len()-1]
  h.data = h.data[:h.Len()-1]
  h.pushDown(0)
  return
}

func (h *Heap) Top() (value H) {
  value = h.data[0]
  return
}

func (h *Heap) Len() int { return len(h.data) }

func (h *Heap) heapify() {
  n := h.Len()
  for i := (n >> 1) - 1; i > -1; i-- {
    h.pushDown(i)
  }
}

func (h *Heap) pushUp(root int) {
  for parent := (root - 1) >> 1; parent >= 0 && h.less(h.data[root], h.data[parent]); parent = (root - 1) >> 1 {
    h.data[root], h.data[parent] = h.data[parent], h.data[root]
    root = parent
  }
}

func (h *Heap) pushDown(root int) {
  n := h.Len()
  for left := (root<<1 + 1); left < n; left = (root<<1 + 1) {
    right := left + 1
    minIndex := root

    if h.less(h.data[left], h.data[minIndex]) {
      minIndex = left
    }

    if right < n && h.less(h.data[right], h.data[minIndex]) {
      minIndex = right
    }

    if minIndex == root {
      return
    }
    h.data[root], h.data[minIndex] = h.data[minIndex], h.data[root]
    root = minIndex
  }
}