package q1382

type TreeNode struct {
	Val   int
	Left  *TreeNode
	Right *TreeNode
}

func flatten(node *TreeNode, result *[]*TreeNode) {
	if node == nil {
		return
	}
	flatten(node.Left, result)
	*result = append(*result, node)
	flatten(node.Right, result)
}

func reconstruct(nodes []*TreeNode) *TreeNode {
	if len(nodes) == 0 {
		return nil
	}
	mid := len(nodes) / 2
	root := nodes[mid]
	root.Left = reconstruct(nodes[:mid])
	root.Right = reconstruct(nodes[mid+1:])

	return root
}

func balanceBST(root *TreeNode) *TreeNode {
	nodes := make([]*TreeNode, 0, 128)
	flatten(root, &nodes)

	return reconstruct(nodes)
}

var _ = balanceBST