1382. Balance a Binary Search Tree

Medium

Given the root of a binary search tree, return a balanced binary search tree with the same node values. If there is more than one answer, return any of them.

A binary search tree is balanced if the depth of the two subtrees of every node never differs by more than 1.

Example 1:

Input: root = [1,null,2,null,3,null,4,null,null]
Output: [2,1,3,null,null,null,4]
Explanation: This is not the only correct answer, [3,1,4,null,2] is also correct.

Example 2:

Input: root = [2,1,3]
Output: [2,1,3]

Constraints:

• The number of nodes in the tree is in the range [1, 10^4].

• 1 <= Node.val <= 10^5

解題

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func balanceBST(root *TreeNode) *TreeNode {
    arr := make([]int, 0)

    var helper func(*TreeNode) 
    helper = func(root *TreeNode) {
        if root == nil { return }
        arr = append(arr, root.Val)

        helper(root.Left)
        helper(root.Right)
    }

    helper(root) // 將樹變成陣列
    sort.Ints(arr)  // 由小至大進行排序

    var combine func([]int) *TreeNode // 陣列左邊 = 比 mid 小，Left child
    combine = func(arr []int) *TreeNode {
        if len(arr) == 0 { return nil }
        if len(arr) == 1 {
            return &TreeNode{ Val: arr[0] }
        }

        mid := len(arr) / 2 
        left := combine(arr[:mid])
        right := combine(arr[mid+1:])
        return &TreeNode{ 
            Val: arr[mid],
            Left : left,
            Right : right,
        }
    }  

    return combine(arr)
}