Balance a Binary Search Tree - Algorithm, Visualization, Examples

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


def inorder(root, arr):
    if root:
        inorder(root.left, arr)
        arr.append(root.val)
        inorder(root.right, arr)


def build_balanced_bst(sorted_arr):
    if not sorted_arr:
        return None
    mid = len(sorted_arr) // 2
    root = TreeNode(sorted_arr[mid])
    root.left = build_balanced_bst(sorted_arr[:mid])
    root.right = build_balanced_bst(sorted_arr[mid+1:])
    return root


def balance_bst(root):
    arr = []
    inorder(root, arr)
    return build_balanced_bst(arr)

# Example usage:
if __name__ == '__main__':
    # Construct an unbalanced BST (right-skewed)
    root = TreeNode(1, None, TreeNode(2, None, TreeNode(3, None, TreeNode(4))))
    balanced = balance_bst(root)
    arr = []
    inorder(balanced, arr)
    print(arr)