58 lines
1.3 KiB
Go
58 lines
1.3 KiB
Go
package q2542
|
|
|
|
import (
|
|
"container/heap"
|
|
"slices"
|
|
)
|
|
|
|
type minHeap struct {
|
|
data []int
|
|
lessFn func(i, j int) bool
|
|
}
|
|
|
|
func (m *minHeap) Len() int { return len(m.data) }
|
|
func (m *minHeap) Less(i int, j int) bool { return m.lessFn(m.data[i], m.data[j]) }
|
|
func (m *minHeap) Push(x any) { m.data = append(m.data, x.(int)) }
|
|
func (m *minHeap) Swap(i int, j int) { m.data[i], m.data[j] = m.data[j], m.data[i] }
|
|
|
|
func (m *minHeap) Pop() any {
|
|
last := m.data[len(m.data)-1]
|
|
m.data = m.data[:len(m.data)-1]
|
|
return last
|
|
}
|
|
|
|
func maxScore(nums1 []int, nums2 []int, k int) int64 {
|
|
indices := make([]int, len(nums1))
|
|
for i := range indices {
|
|
indices[i] = i
|
|
}
|
|
slices.SortFunc(indices, func(a, b int) int { return nums2[b] - nums2[a] })
|
|
|
|
mh := &minHeap{
|
|
data: make([]int, 0, k),
|
|
lessFn: func(a int, b int) bool { return nums1[a] < nums1[b] },
|
|
}
|
|
|
|
sum := 0
|
|
for i := range k {
|
|
mh.data = append(mh.data, indices[i])
|
|
sum += nums1[indices[i]]
|
|
}
|
|
heap.Init(mh)
|
|
maxScore := sum * nums2[indices[k-1]]
|
|
|
|
for i := k; i < len(nums1); i++ {
|
|
realI := indices[i]
|
|
n1, n2 := nums1[realI], nums2[realI]
|
|
if nums1[mh.data[0]] < n1 {
|
|
sum = sum - nums1[mh.data[0]] + n1
|
|
mh.data[0] = realI
|
|
heap.Fix(mh, 0)
|
|
maxScore = max(maxScore, sum*n2)
|
|
}
|
|
}
|
|
|
|
return int64(maxScore)
|
|
}
|
|
|
|
var _ = maxScore
|