// 第 15 周 – 实践课练习的参考答案
// 运行：go run ./w15/exercises
package main

import (
	"container/heap"
	"fmt"
	"math"
	"slices"
)

type Node struct {
	Key         int
	Left, Right *Node
}

func insert(p *Node, x int) *Node {
	if p == nil {
		return &Node{Key: x}
	}
	if x < p.Key {
		p.Left = insert(p.Left, x)
	} else if x > p.Key {
		p.Right = insert(p.Right, x)
	}
	return p
}

// E1. 是否为二叉搜索树？
func isBST(p *Node, lo, hi int) bool {
	if p == nil {
		return true
	}
	if p.Key <= lo || p.Key >= hi {
		return false
	}
	return isBST(p.Left, lo, p.Key) && isBST(p.Right, p.Key, hi)
}

// E2. 第 k 小的元素（中序遍历，提前停止）
func kth(p *Node, k *int) (int, bool) {
	if p == nil {
		return 0, false
	}
	if v, ok := kth(p.Left, k); ok {
		return v, true
	}
	*k--
	if *k == 0 {
		return p.Key, true
	}
	return kth(p.Right, k)
}

// E3. 区间查询
func rangeQuery(p *Node, lo, hi int, out *[]int) {
	if p == nil {
		return
	}
	if lo < p.Key {
		rangeQuery(p.Left, lo, hi, out)
	}
	if lo <= p.Key && p.Key <= hi {
		*out = append(*out, p.Key)
	}
	if p.Key < hi {
		rangeQuery(p.Right, lo, hi, out)
	}
}

// E4. 最近公共祖先
func lca(p *Node, a, b int) int {
	for {
		switch {
		case a < p.Key && b < p.Key:
			p = p.Left
		case a > p.Key && b > p.Key:
			p = p.Right
		default:
			return p.Key
		}
	}
}

// E5. 向下取整（floor）
func floor(p *Node, x int) (int, bool) {
	best, found := 0, false
	for p != nil {
		if p.Key == x {
			return x, true
		}
		if p.Key < x {
			best, found = p.Key, true
			p = p.Right
		} else {
			p = p.Left
		}
	}
	return best, found
}

// E6. 用小顶堆求前 k 大
type MinHeap []int

func (h MinHeap) Len() int           { return len(h) }
func (h MinHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h MinHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *MinHeap) Push(x any)        { *h = append(*h, x.(int)) }
func (h *MinHeap) Pop() any {
	old := *h
	x := old[len(old)-1]
	*h = old[:len(old)-1]
	return x
}

func topK(a []int, k int) []int {
	h := &MinHeap{}
	for _, x := range a {
		heap.Push(h, x)
		if h.Len() > k {
			heap.Pop(h)
		}
	}
	res := slices.Clone(*h)
	slices.Sort(res)
	return res
}

func main() {
	var t *Node
	for _, x := range []int{8, 3, 10, 1, 6, 14, 4, 7, 13} {
		t = insert(t, x)
	}
	fmt.Println("E1. isBST =", isBST(t, math.MinInt, math.MaxInt))
	bad := &Node{8, &Node{3, nil, &Node{Key: 9}}, &Node{Key: 10}}
	fmt.Println("    bad tree isBST =", isBST(bad, math.MinInt, math.MaxInt))
	k := 3
	v, _ := kth(t, &k)
	fmt.Println("E2. the 3rd smallest:", v)
	var r []int
	rangeQuery(t, 5, 11, &r)
	fmt.Println("E3. between [5, 11]:", r)
	fmt.Println("E4. lca(4, 7) =", lca(t, 4, 7), ", lca(1, 13) =", lca(t, 1, 13))
	f1, _ := floor(t, 5)
	f2, _ := floor(t, 12)
	_, ok := floor(t, 0)
	fmt.Println("E5. floor(5) =", f1, ", floor(12) =", f2, ", floor(0) exists?", ok)
	fmt.Println("E6. top 3 [5 1 9 3 7 8 2] =", topK([]int{5, 1, 9, 3, 7, 8, 2}, 3))
}