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