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

English Version

题目描述

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

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

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

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

示例 1：

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

示例 3：

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

提示：

3 <= n <= 10⁵
1 <= nums[i] <= 10⁹
3 <= k <= n
k - 2 <= dist <= n - 2

解法

方法一

Python3

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

Java

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

C++

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

Go

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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

README.md

README.md

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

题目描述

解法

方法一

Python3

Java

C++

Go

Files

README.md

Latest commit

History

README.md

File metadata and controls

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

题目描述

解法

方法一

Python3

Java

C++

Go