#977198 (Go) No.738 平らな農地

提出ソース
結果

問題	No.738 平らな農地
コンテスト
ユーザー	草苺奶昔
提出日時	2024-04-25 00:43:40
言語	Go (1.23.4)
結果	WA
実行時間	-
コード長	6,109 bytes
コンパイル時間	15,236 ms
コンパイル使用メモリ	219,676 KB
実行使用メモリ	12,288 KB
最終ジャッジ日時	2024-11-07 10:09:29
合計ジャッジ時間	23,768 ms
ジャッジサーバーID （参考情報）	judge1 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン	結果
sample	AC * 3 WA * 2
other	AC * 15 WA * 72
権限があれば一括ダウンロードができます
ソースコード

raw source code
// 动态中位数，用两个可删除堆(对顶堆)实现
// 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
}