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