37 lines
639 B
Go
37 lines
639 B
Go
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
|