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

提出ソース
結果

問題	No.738 平らな農地
コンテスト
ユーザー	草苺奶昔
提出日時	2024-04-25 13:36:22
言語	Go (1.23.4)
結果	WA
実行時間	-
コード長	6,181 bytes
コンパイル時間	12,231 ms
コンパイル使用メモリ	234,252 KB
実行使用メモリ	9,600 KB
最終ジャッジ日時	2024-11-08 00:34:44
合計ジャッジ時間	19,787 ms
ジャッジサーバーID （参考情報）	judge1 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン	結果
sample	AC * 4 WA * 1
other	AC * 17 WA * 70
権限があれば一括ダウンロードができます
ソースコード

raw source code
// 动态中位数，用两个可删除堆(对顶堆)实现
// api:
// 1. Insert(x T)
// 2. Discard(x T) bool
// 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()
	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.Println(M.Median())
		// fmt.Println(M.DistToMedian(), 987, M.Size())
	}
	fmt.Fprintln(out, res)
}

type S = int

type DynamicMedian struct {
	size         int32
	lower        *erasableHeap
	upper        *erasableHeap
	lowerCounter map[S]int32
	upperCounter map[S]int32
	lowerSum     S
	upperSum     S
}

func NewDynamicMedian() *DynamicMedian {
	return &DynamicMedian{
		lower:        newErasableHeap(func(a, b S) bool { return a > b }, nil),
		upper:        newErasableHeap(func(a, b S) bool { return a < b }, nil),
		lowerCounter: make(map[S]int32),
		upperCounter: make(map[S]int32),
	}
}

func (d *DynamicMedian) Insert(value S) {
	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) Discard(value S) bool {
	if v1, ok1 := d.lowerCounter[value]; ok1 {
		if v1 == 1 {
			delete(d.lowerCounter, value)
		}
		d.lower.Erase(value)
		d.lowerSum -= value
		d.size--
		d.balance()
		return true
	} else if v2, ok2 := d.upperCounter[value]; ok2 {
		if v2 == 1 {
			delete(d.upperCounter, value)
		}
		d.upper.Erase(value)
		d.upperSum -= value
		d.size--
		d.balance()
		return true
	} else {
		return false
	}
}

// 返回中位数.如果元素个数为偶数,返回两个中位数.
func (d *DynamicMedian) Median() (low, high S) {
	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) DistToMedian() S {
	if d.size == 0 {
		return 0
	}
	low, _ := d.Median()
	sum1 := low*S(d.lower.Len()) - d.lowerSum
	if sum1 < 0 {
		sum1 = -sum1
	}
	sum2 := d.upperSum - low*S(d.upper.Len())
	if sum2 < 0 {
		sum2 = -sum2
	}
	return sum1 + sum2
}

func (d *DynamicMedian) Size() int32 { return d.size }

func (d *DynamicMedian) balance() {
	// 偶数个数时，|lower heap| == |upper heap|
	// 奇数个数时，|lower heap| + 1 == |upper heap|
	for d.lower.Len()+1 < d.upper.Len() {
		upperMin := d.upper.Pop()
		d.lower.Push(upperMin)
		d.lowerSum += upperMin
		d.upperSum -= upperMin
	}
	for d.lower.Len() > d.upper.Len() {
		lowerMin := d.lower.Pop()
		d.upper.Push(lowerMin)
		d.upperSum += lowerMin
		d.lowerSum -= lowerMin
	}

	// if d.size&1 == 0 {
	// 	if d.lower.Len() != d.upper.Len() {
	// 		panic("size error")
	// 	}
	// } else {
	// 	if d.lower.Len()+1 != d.upper.Len() {
	// 		panic("size error")
	// 	}
	// }

	if d.lower.Len() == 0 || d.upper.Len() == 0 {
		return
	}

	if d.lower.Peek() > d.upper.Peek() {
		upperMin := d.upper.Pop()
		d.lower.Push(upperMin)
		d.lowerSum += upperMin
		d.upperSum -= upperMin

		lowerMax := d.lower.Pop()
		d.upper.Push(lowerMax)
		d.upperSum += lowerMax
		d.lowerSum -= lowerMax
	}
}

type erasableHeap struct {
	data   *heap
	erased *heap
	size   int32
}

func newErasableHeap(less func(a, b S) bool, nums []S) *erasableHeap {
	return &erasableHeap{newHeap(less, nums), newHeap(less, nil), 0}
}

// 从堆中删除一个元素,要保证堆中存在该元素.
func (h *erasableHeap) Erase(value S) {
	h.erased.Push(value)
	h.normalize()
	h.size--
}

func (h *erasableHeap) Push(value S) {
	h.data.Push(value)
	h.normalize()
	h.size++
}

func (h *erasableHeap) Pop() (value S) {
	value = h.data.Pop()
	h.normalize()
	h.size--
	return
}

func (h *erasableHeap) Peek() (value S) {
	value = h.data.Top()
	return
}

func (h *erasableHeap) Len() int32 {
	return h.size
}

func (h *erasableHeap) Clear() {
	h.data.Clear()
	h.erased.Clear()
	h.size = 0
}

func (h *erasableHeap) normalize() {
	for h.data.Len() > 0 && h.erased.Len() > 0 && h.data.Top() == h.erased.Top() {
		h.data.Pop()
		h.erased.Pop()
	}
}

func newHeap(less func(a, b S) bool, nums []S) *heap {
	nums = append(nums[:0:0], nums...)
	heap := &heap{less: less, data: nums}
	if len(nums) > 1 {
		heap.heapify()
	}
	return heap
}

type heap struct {
	data []S
	less func(a, b S) bool
}

func (h *heap) Push(value S) {
	h.data = append(h.data, value)
	h.pushUp(h.Len() - 1)
}

func (h *heap) Pop() (value S) {
	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 *heap) Top() (value S) {
	if h.Len() == 0 {
		panic("heap is empty")
	}
	value = h.data[0]
	return
}

func (h *heap) Len() int { return len(h.data) }

func (h *heap) Clear() {
	h.data = h.data[:0]
}

func (h *heap) heapify() {
	n := h.Len()
	for i := (n >> 1) - 1; i > -1; i-- {
		h.pushDown(i)
	}
}

func (h *heap) 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 *heap) 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
}