// 动态中位数,用两个可删除堆(对顶堆)实现 // api: // 1. Insert(x T) // 2. Erase(x T) // 3. Median() (low, high T) // 4. DistToMedian() T // 5. Size() int32 package main import ( "bufio" "fmt" "os" ) func main() { yuki738() } const INF int = 1e18 // No.738 平らな農地 // https://yukicoder.me/problems/no/738 // !滑动窗口所有数到中位数的距离和 func yuki738() { in := bufio.NewReader(os.Stdin) out := bufio.NewWriter(os.Stdout) defer out.Flush() var n, k int32 fmt.Fscan(in, &n, &k) nums := make([]int, n) for i := int32(0); i < n; i++ { fmt.Fscan(in, &nums[i]) } M := NewDynamicMedian[int]() res := INF for i := int32(0); i < n; i++ { M.Insert(nums[i]) if i >= k { M.Discard(nums[i-k]) } if i >= k-1 { res = min(res, M.DistToMedian()) } } fmt.Fprintln(out, res) } type Integer interface { int | uint | int8 | uint8 | int16 | uint16 | int32 | uint32 | int64 | uint64 } type DynamicMedian[V Integer] struct { size int32 lower *erasableHeapGeneric[V] upper *erasableHeapGeneric[V] lowerSum V upperSum V lowerCounter map[V]int32 upperCounter map[V]int32 } func NewDynamicMedian[V Integer]() *DynamicMedian[V] { return &DynamicMedian[V]{ lower: newErasableHeapGeneric[V](func(a, b V) bool { return a > b }), upper: newErasableHeapGeneric[V](func(a, b V) bool { return a < b }), lowerCounter: make(map[V]int32), upperCounter: make(map[V]int32), } } func (d *DynamicMedian[V]) Insert(value V) { if d.size&1 == 0 { d.upper.Push(value) d.upperSum += value d.upperCounter[value]++ } else { d.lower.Push(value) d.lowerSum += value d.lowerCounter[value]++ } d.size++ d.balance() } func (d *DynamicMedian[V]) Discard(value V) bool { if d.lowerCounter[value] > 0 { d.lowerCounter[value]-- d.lower.Erase(value) d.lowerSum -= value d.size-- return true } else if d.upperCounter[value] > 0 { d.upperCounter[value]-- d.upper.Erase(value) d.upperSum -= value d.size-- return true } else { return false } } // 返回中位数.如果元素个数为偶数,返回两个中位数. func (d *DynamicMedian[V]) Median() (low, high V) { if d.size == 0 { return } if d.size&1 == 0 { low = d.lower.Peek() high = d.upper.Peek() } else { low = d.upper.Peek() high = low } return } func (d *DynamicMedian[V]) DistToMedian() V { if d.size == 0 { return 0 } low, _ := d.Median() sum1 := low*V(d.lower.Len()) - d.lowerSum if sum1 < 0 { sum1 = -sum1 } sum2 := low*V(d.upper.Len()) - d.upperSum if sum2 < 0 { sum2 = -sum2 } return sum1 + sum2 } func (d *DynamicMedian[V]) Size() int32 { return d.size } func (d *DynamicMedian[V]) balance() { // 偶数个数时,|lower heap| == |upper heap| // 奇数个数时,|lower heap| + 1 == |upper heap| for d.lower.Len()+1 < d.upper.Len() { d.lower.Push(d.upper.Pop()) } for d.lower.Len() > d.upper.Len() { d.upper.Push(d.lower.Pop()) } if d.lower.Len() == 0 || d.upper.Len() == 0 { return } if v1, v2 := d.lower.Peek(), d.upper.Peek(); v1 > v2 { upperValue := d.upper.Pop() d.lower.Push(upperValue) d.lowerSum += upperValue d.upperSum -= upperValue lowerValue := d.lower.Pop() d.upper.Push(lowerValue) d.upperSum += lowerValue d.lowerSum -= lowerValue } } type erasableHeapGeneric[V comparable] struct { data *heapGeneric[V] erased *heapGeneric[V] } func newErasableHeapGeneric[V comparable](less func(a, b V) bool, nums ...V) *erasableHeapGeneric[V] { return &erasableHeapGeneric[V]{newHeapGeneric(less, nums...), newHeapGeneric(less)} } // 从堆中删除一个元素,要保证堆中存在该元素. func (h *erasableHeapGeneric[V]) Erase(value V) { h.erased.Push(value) h.normalize() } func (h *erasableHeapGeneric[V]) Push(value V) { h.data.Push(value) h.normalize() } func (h *erasableHeapGeneric[V]) Pop() (value V) { value = h.data.Pop() h.normalize() return } func (h *erasableHeapGeneric[V]) Peek() (value V) { value = h.data.Top() return } func (h *erasableHeapGeneric[V]) Len() int { return h.data.Len() } func (h *erasableHeapGeneric[V]) Clear() { h.data.Clear() h.erased.Clear() } func (h *erasableHeapGeneric[V]) normalize() { for h.data.Len() > 0 && h.erased.Len() > 0 && h.data.Top() == h.erased.Top() { h.data.Pop() h.erased.Pop() } } type heapGeneric[V comparable] struct { data []V less func(a, b V) bool } func newHeapGeneric[V comparable](less func(a, b V) bool, nums ...V) *heapGeneric[V] { nums = append(nums[:0:0], nums...) heap := &heapGeneric[V]{less: less, data: nums} if len(nums) > 1 { heap.heapify() } return heap } func (h *heapGeneric[V]) Push(value V) { h.data = append(h.data, value) h.pushUp(h.Len() - 1) } func (h *heapGeneric[V]) Pop() (value V) { if h.Len() == 0 { panic("heap is empty") } value = h.data[0] h.data[0] = h.data[h.Len()-1] h.data = h.data[:h.Len()-1] h.pushDown(0) return } func (h *heapGeneric[V]) Top() (value V) { if h.Len() == 0 { panic("heap is empty") } value = h.data[0] return } func (h *heapGeneric[V]) Len() int { return len(h.data) } func (h *heapGeneric[V]) Clear() { h.data = h.data[:0] } func (h *heapGeneric[V]) heapify() { n := h.Len() for i := (n >> 1) - 1; i > -1; i-- { h.pushDown(i) } } func (h *heapGeneric[V]) pushUp(root int) { for parent := (root - 1) >> 1; parent >= 0 && h.less(h.data[root], h.data[parent]); parent = (root - 1) >> 1 { h.data[root], h.data[parent] = h.data[parent], h.data[root] root = parent } } func (h *heapGeneric[V]) pushDown(root int) { n := h.Len() for left := (root<<1 + 1); left < n; left = (root<<1 + 1) { right := left + 1 minIndex := root if h.less(h.data[left], h.data[minIndex]) { minIndex = left } if right < n && h.less(h.data[right], h.data[minIndex]) { minIndex = right } if minIndex == root { return } h.data[root], h.data[minIndex] = h.data[minIndex], h.data[root] root = minIndex } } func min(a, b int) int { if a < b { return a } return b } func max(a, b int) int { if a > b { return a } return b }