1382. Balance a Binary Search Tree - cocoder39/coco39_LC GitHub Wiki

1382. Balance a Binary Search Tree

• Step 1: in-order traverse to get sorted order
• step 2: divide-and-conquer to build balanced BST

class Solution:
    def balanceBST(self, root: TreeNode) -> TreeNode:
        sortedArr = []
        self.inorderTraverse(root, sortedArr)
        
        return self.formBalancedBST(sortedArr, 0, len(sortedArr)-1)
    
    def inorderTraverse(self, root: TreeNode, sortedArr: list) -> None:
        if not root:
            return 
        
        if root.left:
            self.inorderTraverse(root.left, sortedArr)
            
        sortedArr.append(root)
        
        if root.right:
            self.inorderTraverse(root.right, sortedArr)
    
    def formBalancedBST(self, sortedArr: list, start: int, end: int) -> TreeNode:
        if start > end:
            return None
        
        mid = start + (end - start) // 2
        node = sortedArr[mid]
        node.left = self.formBalancedBST(sortedArr, start, mid-1)
        node.right = self.formBalancedBST(sortedArr, mid+1, end)
        return node