3013. Divide an Array Into Subarrays With Minimum Cost II

Description

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.

 

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 <= 105
  • 1 <= nums[i] <= 109
  • 3 <= k <= n
  • k - 2 <= dist <= n - 2

Solutions

Solution 1

Python Code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

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 Code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

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++ Code
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

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 Code
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194

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