package top100liked

import (
	"fmt"
	"testing"
)

// https://leetcode.cn/problems/find-median-from-data-stream/?envType=study-plan-v2&envId=top-100-liked

type heap[T any] struct {
	data        []T
	compareFunc func(T, T) bool
}

func NewHeap[T any](compareFunc func(T, T) bool) heap[T] {
	return heap[T]{
		compareFunc: compareFunc,
	}
}

func (h *heap[T]) Up(idx int) {
	for (idx-1)/2 >= 0 {
		cur := h.data[idx]
		p := h.data[(idx-1)/2]
		if h.compareFunc(cur, p) {
			h.data[idx], h.data[(idx-1)/2] = h.data[(idx-1)/2], h.data[idx]
			idx = (idx - 1) / 2
		} else {
			break
		}
	}
}

func (h *heap[T]) Down(idx int) {
	for {
		l, r := 2*idx+1, 2*idx+2
		m := idx
		if l < len(h.data) && h.compareFunc(h.data[l], h.data[m]) {
			m = l
		}
		if r < len(h.data) && h.compareFunc(h.data[r], h.data[m]) {
			m = r
		}
		if m == idx {
			break
		}
		h.data[idx], h.data[m] = h.data[m], h.data[idx]
		idx = m
	}
}

func (h *heap[T]) Push(val T) {
	h.data = append(h.data, val)
	h.Up(len(h.data) - 1)
}

func (h *heap[T]) Pop() T {
	if len(h.data) <= 0 {
		var zero T
		return zero
	}
	old := h.data[0]
	h.data[0] = h.data[len(h.data)-1]
	h.data = h.data[:len(h.data)-1]
	h.Down(0)
	return old
}

func (h *heap[T]) Top() T {
	if len(h.data) > 0 {
		return h.data[0]
	}
	var zero T
	return zero
}

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

type MedianFinder struct {
	left  heap[int] // 小
	right heap[int] // 大
}

func Constructor() MedianFinder {
	return MedianFinder{
		left:  NewHeap(func(i1, i2 int) bool { return i1 > i2 }),
		right: NewHeap(func(i1, i2 int) bool { return i1 < i2 }),
	}
}

func (this *MedianFinder) AddNum(num int) {
	if this.left.Len() == 0 {
		this.left.Push(num)
		return
	}
	if num >= this.left.Top() {
		this.right.Push(num)
	} else {
		this.left.Push(num)
	}
	if this.left.Len() > this.right.Len()+1 {
		v := this.left.Pop()
		this.right.Push(v)
	}
	if this.right.Len() > this.left.Len()+1 {
		v := this.right.Pop()
		this.left.Push(v)
	}
}

func (this *MedianFinder) FindMedian() float64 {
	switch (this.left.Len() + this.right.Len()) % 2 {
	case 0:
		return (float64(this.left.Top()) + float64(this.right.Top())) / 2
	case 1:
		if this.left.Len() > this.right.Len() {
			return float64(this.left.Top())
		} else {
			return float64(this.right.Top())
		}
	}
	return 0
}

func Test(t *testing.T) {
	medianFinder := Constructor()
	medianFinder.AddNum(1)                 // arr = [1]
	medianFinder.AddNum(2)                 // arr = [1, 2]
	fmt.Println(medianFinder.FindMedian()) // 返回 1.5 ((1 + 2) / 2)
	medianFinder.AddNum(3)                 // arr[1, 2, 3]
	fmt.Println(medianFinder.FindMedian()) // return 2.0
}