class Solution {
    public static class TreeNode {
        public int val;
        TreeNode left;
        TreeNode right;
        public TreeNode(){}
        public TreeNode(int val) {this.val = val;}
        public TreeNode(int val, TreeNode left, TreeNode right) {
            this.val = val;
            this.left  =left;
            this.right = right;
        }
    }

    public List<Integer> nodes = new ArrayList<Integer>();

    public void traversal(TreeNode root) {
        if (root == null) return;
        traversal(root.left);
        nodes.add(root.val);
        traversal(root.right);
    }

    public TreeNode BST(int left, int right) {
        if (left > right) return null;
        int mid = left + (right - left) / 2;

        TreeNode current = new TreeNode(nodes.get(mid));
        current.left = BST(left, mid - 1);
        current.right = BST(mid + 1, right);

        return current;
    }

    public TreeNode balanceBST(TreeNode root) {
        traversal(root);
        return BST(0, nodes.size() - 1);
    }
}