62 lines
1.2 KiB
Go
62 lines
1.2 KiB
Go
package q2462
|
|
|
|
import "container/heap"
|
|
|
|
type CustomMinHeap struct {
|
|
data []int
|
|
less func(a, b int) bool
|
|
}
|
|
|
|
func (c *CustomMinHeap) Len() int { return len(c.data) }
|
|
func (c *CustomMinHeap) Less(i int, j int) bool { return c.less(c.data[i], c.data[j]) }
|
|
func (c *CustomMinHeap) Push(x any) { c.data = append(c.data, x.(int)) }
|
|
func (c *CustomMinHeap) Swap(i int, j int) { c.data[i], c.data[j] = c.data[j], c.data[i] }
|
|
|
|
func (c *CustomMinHeap) Pop() any {
|
|
last := c.data[len(c.data)-1]
|
|
c.data = c.data[:len(c.data)-1]
|
|
return last
|
|
}
|
|
|
|
func totalCost(costs []int, k int, candidates int) int64 {
|
|
lessFn := func(i, j int) bool {
|
|
if costs[i] == costs[j] {
|
|
return i < j
|
|
}
|
|
return costs[i] < costs[j]
|
|
}
|
|
|
|
indices := make([]int, 0, 2*candidates)
|
|
l, r := 0, len(costs)-1
|
|
for l <= r && l < candidates {
|
|
indices = append(indices, l)
|
|
if l < r {
|
|
indices = append(indices, r)
|
|
}
|
|
l++
|
|
r--
|
|
}
|
|
|
|
hp := &CustomMinHeap{data: indices, less: lessFn}
|
|
heap.Init(hp)
|
|
|
|
cost := 0
|
|
for range k {
|
|
i := heap.Pop(hp).(int)
|
|
cost += costs[i]
|
|
|
|
if l <= r {
|
|
switch {
|
|
case i < l:
|
|
heap.Push(hp, l)
|
|
l++
|
|
default:
|
|
heap.Push(hp, r)
|
|
r--
|
|
}
|
|
}
|
|
}
|
|
return int64(cost)
|
|
}
|
|
|
|
var _ = totalCost
|