|
63 | 63 | <!-- tabs:start -->
|
64 | 64 |
|
65 | 65 | ```python
|
| 66 | +from sortedcontainers import SortedList |
66 | 67 |
|
| 68 | + |
| 69 | +class Solution: |
| 70 | + def minimumCost(self, nums: List[int], k: int, dist: int) -> int: |
| 71 | + n = len(nums) |
| 72 | + |
| 73 | + sl = SortedList() |
| 74 | + y = nums[0] |
| 75 | + ans = float("inf") |
| 76 | + i = 1 |
| 77 | + running_sum = 0 |
| 78 | + |
| 79 | + for j in range(1, n): |
| 80 | + pos = bisect.bisect_left(sl, nums[j]) |
| 81 | + sl.add(nums[j]) |
| 82 | + |
| 83 | + if pos < k - 1: |
| 84 | + running_sum += nums[j] |
| 85 | + if len(sl) > k - 1: |
| 86 | + running_sum -= sl[k - 1] |
| 87 | + |
| 88 | + while j - i > dist: |
| 89 | + removed_pos = sl.index(nums[i]) |
| 90 | + removed_element = nums[i] |
| 91 | + sl.remove(removed_element) |
| 92 | + |
| 93 | + if removed_pos < k - 1: |
| 94 | + running_sum -= removed_element |
| 95 | + if len(sl) >= k - 1: |
| 96 | + running_sum += sl[k - 2] |
| 97 | + i += 1 |
| 98 | + |
| 99 | + if j - i + 1 >= k - 1: |
| 100 | + ans = min(ans, running_sum) |
| 101 | + |
| 102 | + return ans + y |
67 | 103 | ```
|
68 | 104 |
|
69 | 105 | ```java
|
70 |
| - |
| 106 | +class Solution { |
| 107 | + public long minimumCost(int[] nums, int k, int dist) { |
| 108 | + long result = Long.MAX_VALUE, sum = 0L; |
| 109 | + int n = nums.length; |
| 110 | + TreeSet<Integer> set1 |
| 111 | + = new TreeSet<>((a, b) -> nums[a] == nums[b] ? a - b : nums[a] - nums[b]); |
| 112 | + TreeSet<Integer> set2 |
| 113 | + = new TreeSet<>((a, b) -> nums[a] == nums[b] ? a - b : nums[a] - nums[b]); |
| 114 | + for (int i = 1; i < n; i++) { |
| 115 | + set1.add(i); |
| 116 | + sum += nums[i]; |
| 117 | + if (set1.size() >= k) { |
| 118 | + int x = set1.pollLast(); |
| 119 | + sum -= nums[x]; |
| 120 | + set2.add(x); |
| 121 | + } |
| 122 | + if (i - dist > 0) { |
| 123 | + result = Math.min(result, sum); |
| 124 | + int temp = i - dist; |
| 125 | + if (set1.contains(temp)) { |
| 126 | + set1.remove(temp); |
| 127 | + sum -= nums[temp]; |
| 128 | + if (set2.size() > 0) { |
| 129 | + int y = set2.pollFirst(); |
| 130 | + sum += nums[y]; |
| 131 | + set1.add(y); |
| 132 | + } |
| 133 | + } else { |
| 134 | + set2.remove(i - dist); |
| 135 | + } |
| 136 | + } |
| 137 | + } |
| 138 | + return result + nums[0]; |
| 139 | + } |
| 140 | +} |
71 | 141 | ```
|
72 | 142 |
|
73 | 143 | ```cpp
|
| 144 | +class Solution { |
| 145 | +public: |
| 146 | + long long minimumCost(vector<int>& nums, int k, int dist) { |
| 147 | + multiset<int> sml, big; |
| 148 | + int sz = dist + 1; |
| 149 | + long long sum = 0, ans = 0; |
| 150 | + for (int i = 1; i <= sz; i++) { |
| 151 | + sml.insert(nums[i]); |
| 152 | + sum += nums[i]; |
| 153 | + } |
| 154 | + while (sml.size() > k - 1) { |
| 155 | + big.insert(*sml.rbegin()); |
| 156 | + sum -= *sml.rbegin(); |
| 157 | + sml.erase(sml.find(*sml.rbegin())); |
| 158 | + } |
| 159 | + ans = sum; |
| 160 | + for (int i = sz + 1; i < nums.size(); i++) { |
| 161 | + sum += nums[i]; |
| 162 | + sml.insert(nums[i]); |
| 163 | + if (big.find(nums[i - sz]) != big.end()) { |
| 164 | + big.erase(big.find(nums[i - sz])); |
| 165 | + } else { |
| 166 | + sum -= nums[i - sz]; |
| 167 | + sml.erase(sml.find(nums[i - sz])); |
| 168 | + } |
74 | 169 |
|
| 170 | + while (sml.size() > k - 1) { |
| 171 | + sum -= *sml.rbegin(); |
| 172 | + big.insert(*sml.rbegin()); |
| 173 | + sml.erase(sml.find(*sml.rbegin())); |
| 174 | + } |
| 175 | + while (sml.size() < k - 1) { |
| 176 | + sum += *big.begin(); |
| 177 | + sml.insert(*big.begin()); |
| 178 | + big.erase(big.begin()); |
| 179 | + } |
| 180 | + while (!sml.empty() && !big.empty() && *sml.rbegin() > *big.begin()) { |
| 181 | + sum -= *sml.rbegin() - *big.begin(); |
| 182 | + sml.insert(*big.begin()); |
| 183 | + big.insert(*sml.rbegin()); |
| 184 | + sml.erase(sml.find(*sml.rbegin())); |
| 185 | + big.erase(big.begin()); |
| 186 | + } |
| 187 | + ans = min(ans, sum); |
| 188 | + } |
| 189 | + int p = 0; |
| 190 | + return nums[0] + ans; |
| 191 | + } |
| 192 | +}; |
75 | 193 | ```
|
76 | 194 |
|
77 | 195 | ```go
|
| 196 | +func minimumCost(nums []int, k int, dist int) int64 { |
| 197 | + res := nums[0] + slices.Min(windowTopKSum(nums[1:], dist+1, k-1, true)) |
| 198 | + return int64(res) |
| 199 | +} |
| 200 | + |
| 201 | +func windowTopKSum(nums []int, windowSize, k int, min bool) []int { |
| 202 | + n := len(nums) |
| 203 | + ts := NewTopKSum(k, min) |
| 204 | + res := []int{} |
| 205 | + for right := 0; right < n; right++ { |
| 206 | + ts.Add(nums[right]) |
| 207 | + if right >= windowSize { |
| 208 | + ts.Discard(nums[right-windowSize]) |
| 209 | + } |
| 210 | + if right >= windowSize-1 { |
| 211 | + res = append(res, ts.Query()) |
| 212 | + } |
| 213 | + } |
| 214 | + return res |
| 215 | +} |
| 216 | + |
| 217 | +type TopKSum struct { |
| 218 | + sum int |
| 219 | + k int |
| 220 | + in *Heap |
| 221 | + out *Heap |
| 222 | + dIn *Heap |
| 223 | + dOut *Heap |
| 224 | + counter map[int]int |
| 225 | +} |
| 226 | + |
| 227 | +func NewTopKSum(k int, min bool) *TopKSum { |
| 228 | + var less func(a, b int) bool |
| 229 | + if min { |
| 230 | + less = func(a, b int) bool { return a < b } |
| 231 | + } else { |
| 232 | + less = func(a, b int) bool { return a > b } |
| 233 | + } |
| 234 | + return &TopKSum{ |
| 235 | + k: k, |
| 236 | + in: NewHeap(less), |
| 237 | + out: NewHeap(less), |
| 238 | + dIn: NewHeap(less), |
| 239 | + dOut: NewHeap(less), |
| 240 | + counter: map[int]int{}, |
| 241 | + } |
| 242 | +} |
| 243 | + |
| 244 | +func (t *TopKSum) Query() int { |
| 245 | + return t.sum |
| 246 | +} |
| 247 | + |
| 248 | +func (t *TopKSum) Add(x int) { |
| 249 | + t.counter[x]++ |
| 250 | + t.in.Push(-x) |
| 251 | + t.sum += x |
| 252 | + t.modify() |
| 253 | +} |
| 254 | + |
| 255 | +func (t *TopKSum) Discard(x int) bool { |
| 256 | + if t.counter[x] == 0 { |
| 257 | + return false |
| 258 | + } |
| 259 | + t.counter[x]-- |
| 260 | + if t.in.Len() > 0 && -t.in.Top() == x { |
| 261 | + t.sum -= x |
| 262 | + t.in.Pop() |
| 263 | + } else if t.in.Len() > 0 && -t.in.Top() > x { |
| 264 | + t.sum -= x |
| 265 | + t.dIn.Push(-x) |
| 266 | + } else { |
| 267 | + t.dOut.Push(x) |
| 268 | + } |
| 269 | + t.modify() |
| 270 | + return true |
| 271 | +} |
| 272 | + |
| 273 | +func (t *TopKSum) SetK(k int) { |
| 274 | + t.k = k |
| 275 | + t.modify() |
| 276 | +} |
| 277 | + |
| 278 | +func (t *TopKSum) GetK() int { |
| 279 | + return t.k |
| 280 | +} |
| 281 | + |
| 282 | +func (t *TopKSum) Len() int { |
| 283 | + return t.in.Len() + t.out.Len() - t.dIn.Len() - t.dOut.Len() |
| 284 | +} |
| 285 | + |
| 286 | +func (t *TopKSum) Has(x int) bool { |
| 287 | + return t.counter[x] > 0 |
| 288 | +} |
| 289 | + |
| 290 | +func (t *TopKSum) modify() { |
| 291 | + for t.out.Len() > 0 && (t.in.Len()-t.dIn.Len() < t.k) { |
| 292 | + p := t.out.Pop() |
| 293 | + if t.dOut.Len() > 0 && p == t.dOut.Top() { |
| 294 | + t.dOut.Pop() |
| 295 | + } else { |
| 296 | + t.sum += p |
| 297 | + t.in.Push(-p) |
| 298 | + } |
| 299 | + } |
| 300 | + |
| 301 | + for t.in.Len()-t.dIn.Len() > t.k { |
| 302 | + p := -t.in.Pop() |
| 303 | + if t.dIn.Len() > 0 && p == -t.dIn.Top() { |
| 304 | + t.dIn.Pop() |
| 305 | + } else { |
| 306 | + t.sum -= p |
| 307 | + t.out.Push(p) |
| 308 | + } |
| 309 | + } |
| 310 | + |
| 311 | + for t.dIn.Len() > 0 && t.in.Top() == t.dIn.Top() { |
| 312 | + t.in.Pop() |
| 313 | + t.dIn.Pop() |
| 314 | + } |
| 315 | +} |
| 316 | + |
| 317 | +type H = int |
| 318 | + |
| 319 | +func NewHeap(less func(a, b H) bool, nums ...H) *Heap { |
| 320 | + nums = append(nums[:0:0], nums...) |
| 321 | + heap := &Heap{less: less, data: nums} |
| 322 | + heap.heapify() |
| 323 | + return heap |
| 324 | +} |
| 325 | + |
| 326 | +type Heap struct { |
| 327 | + data []H |
| 328 | + less func(a, b H) bool |
| 329 | +} |
| 330 | + |
| 331 | +func (h *Heap) Push(value H) { |
| 332 | + h.data = append(h.data, value) |
| 333 | + h.pushUp(h.Len() - 1) |
| 334 | +} |
| 335 | + |
| 336 | +func (h *Heap) Pop() (value H) { |
| 337 | + if h.Len() == 0 { |
| 338 | + panic("heap is empty") |
| 339 | + } |
| 340 | + |
| 341 | + value = h.data[0] |
| 342 | + h.data[0] = h.data[h.Len()-1] |
| 343 | + h.data = h.data[:h.Len()-1] |
| 344 | + h.pushDown(0) |
| 345 | + return |
| 346 | +} |
| 347 | + |
| 348 | +func (h *Heap) Top() (value H) { |
| 349 | + value = h.data[0] |
| 350 | + return |
| 351 | +} |
| 352 | + |
| 353 | +func (h *Heap) Len() int { return len(h.data) } |
| 354 | + |
| 355 | +func (h *Heap) heapify() { |
| 356 | + n := h.Len() |
| 357 | + for i := (n >> 1) - 1; i > -1; i-- { |
| 358 | + h.pushDown(i) |
| 359 | + } |
| 360 | +} |
| 361 | + |
| 362 | +func (h *Heap) pushUp(root int) { |
| 363 | + for parent := (root - 1) >> 1; parent >= 0 && h.less(h.data[root], h.data[parent]); parent = (root - 1) >> 1 { |
| 364 | + h.data[root], h.data[parent] = h.data[parent], h.data[root] |
| 365 | + root = parent |
| 366 | + } |
| 367 | +} |
| 368 | + |
| 369 | +func (h *Heap) pushDown(root int) { |
| 370 | + n := h.Len() |
| 371 | + for left := (root<<1 + 1); left < n; left = (root<<1 + 1) { |
| 372 | + right := left + 1 |
| 373 | + minIndex := root |
| 374 | + |
| 375 | + if h.less(h.data[left], h.data[minIndex]) { |
| 376 | + minIndex = left |
| 377 | + } |
| 378 | + |
| 379 | + if right < n && h.less(h.data[right], h.data[minIndex]) { |
| 380 | + minIndex = right |
| 381 | + } |
78 | 382 |
|
| 383 | + if minIndex == root { |
| 384 | + return |
| 385 | + } |
| 386 | + h.data[root], h.data[minIndex] = h.data[minIndex], h.data[root] |
| 387 | + root = minIndex |
| 388 | + } |
| 389 | +} |
79 | 390 | ```
|
80 | 391 |
|
81 | 392 | <!-- tabs:end -->
|
|
0 commit comments