115 lines
2.5 KiB
Go
115 lines
2.5 KiB
Go
package q3013
|
|
|
|
import (
|
|
"container/heap"
|
|
"math"
|
|
)
|
|
|
|
type HeapElem struct {
|
|
pos int
|
|
heapPos int
|
|
}
|
|
|
|
type Heap struct {
|
|
data []int
|
|
h []*HeapElem
|
|
isMin bool
|
|
}
|
|
|
|
func (h *Heap) Len() int { return len(h.h) }
|
|
func (h *Heap) TopVal() int { return h.data[h.h[0].pos] }
|
|
|
|
func (h *Heap) Less(i int, j int) bool {
|
|
a, b := h.data[h.h[i].pos], h.data[h.h[j].pos]
|
|
if h.isMin {
|
|
return a < b
|
|
}
|
|
return a > b
|
|
}
|
|
|
|
func (h *Heap) Push(x any) {
|
|
elem := x.(*HeapElem)
|
|
elem.heapPos = len(h.h)
|
|
h.h = append(h.h, elem)
|
|
}
|
|
|
|
func (h *Heap) Pop() any {
|
|
ret := h.h[len(h.h)-1]
|
|
h.h = h.h[:len(h.h)-1]
|
|
ret.heapPos = -1
|
|
return ret
|
|
}
|
|
|
|
func (h *Heap) Swap(i int, j int) {
|
|
h.h[i].heapPos = j
|
|
h.h[j].heapPos = i
|
|
h.h[i], h.h[j] = h.h[j], h.h[i]
|
|
}
|
|
|
|
var (
|
|
minHpElems = make([]HeapElem, 100_000)
|
|
maxHpElems = make([]HeapElem, 100_000)
|
|
minHeapBuf = make([]*HeapElem, 100_000)
|
|
maxHeapBuf = make([]*HeapElem, 100_000)
|
|
)
|
|
|
|
func minimumCost(nums []int, k int, dist int) int64 {
|
|
minHpElems = minHpElems[:0]
|
|
maxHpElems = maxHpElems[:0]
|
|
for i := range nums {
|
|
minHpElems = append(minHpElems, HeapElem{pos: i, heapPos: -1})
|
|
maxHpElems = append(maxHpElems, HeapElem{pos: i, heapPos: -1})
|
|
}
|
|
maxHeap := Heap{data: nums, h: maxHeapBuf[:0], isMin: false}
|
|
minHeap := Heap{data: nums, h: minHeapBuf[:0], isMin: true}
|
|
|
|
sum := nums[0]
|
|
minSum := math.MaxInt
|
|
|
|
for i := 1; i < len(nums); i++ {
|
|
remIdx := i - dist - 1
|
|
if remIdx > 0 {
|
|
if minHpElems[remIdx].heapPos != -1 {
|
|
heap.Remove(&minHeap, minHpElems[remIdx].heapPos)
|
|
} else if maxHpElems[remIdx].heapPos != -1 {
|
|
heap.Remove(&maxHeap, maxHpElems[remIdx].heapPos)
|
|
sum -= nums[remIdx]
|
|
}
|
|
}
|
|
|
|
if maxHeap.Len() < k-1 {
|
|
if minHeap.Len() == 0 || nums[i] <= minHeap.TopVal() {
|
|
heap.Push(&maxHeap, &maxHpElems[i])
|
|
sum += nums[i]
|
|
} else {
|
|
heap.Push(&maxHeap, &maxHpElems[minHeap.h[0].pos])
|
|
sum += nums[minHeap.h[0].pos]
|
|
minHeap.h[0].heapPos = -1
|
|
minHpElems[i].heapPos = 0
|
|
minHeap.h[0] = &minHpElems[i]
|
|
heap.Fix(&minHeap, 0)
|
|
}
|
|
|
|
} else {
|
|
heap.Push(&minHeap, &minHpElems[i])
|
|
minTopV, maxTopV := minHeap.TopVal(), maxHeap.TopVal()
|
|
if minTopV < maxTopV {
|
|
sum += minTopV - maxTopV
|
|
|
|
minHeap.h[0].heapPos, maxHeap.h[0].heapPos = -1, -1
|
|
minHeap.h[0], maxHeap.h[0] = &minHpElems[maxHeap.h[0].pos], &maxHpElems[minHeap.h[0].pos]
|
|
minHeap.h[0].heapPos, maxHeap.h[0].heapPos = 0, 0
|
|
heap.Fix(&minHeap, 0)
|
|
heap.Fix(&maxHeap, 0)
|
|
}
|
|
}
|
|
|
|
if maxHeap.Len() == k-1 {
|
|
minSum = min(minSum, sum)
|
|
}
|
|
}
|
|
|
|
return int64(minSum)
|
|
}
|
|
|
|
var _ = minimumCost
|