Update Balance a Binary Search Tree.java

Author: varunu28

Medium/Balance a Binary Search Tree.java

@@ -14,29 +14,37 @@
  * }
  */
 class Solution {
-  public TreeNode balanceBST(TreeNode root) {
-    List<Integer> sortedList = new ArrayList<>();
-    inorderTraversal(root, sortedList);
-    return buildBST(sortedList, 0, sortedList.size() - 1);
-  }   
-  
-  private TreeNode buildBST(List<Integer> sortedList, int start, int end) {
-    if (start > end) {
-      return null;
+    public TreeNode balanceBST(TreeNode root) {
+        List<Integer> inorder = getInorderNodes(root);
+        return createBST(inorder, 0, inorder.size() - 1);
     }
-    int mid = (start + end) / 2;
-    TreeNode root = new TreeNode(sortedList.get(mid));
-    root.left = buildBST(sortedList, start, mid - 1);
-    root.right = buildBST(sortedList, mid + 1, end);
-    return root;
-  }
-  
-  private void inorderTraversal(TreeNode root, List<Integer> sortedList) {
-    if (root == null) {
-      return;
+
+    private List<Integer> getInorderNodes(TreeNode root) {
+        Stack<TreeNode> stack = new Stack<>();
+        while (root != null) {
+            stack.push(root);
+            root = root.left;
+        }
+        List<Integer> list = new ArrayList<>();
+        while (!stack.isEmpty()) {
+            TreeNode node = stack.pop();
+            list.add(node.val);
+            TreeNode rightNode = node.right;
+            while (rightNode != null) {
+                stack.push(rightNode);
+                rightNode = rightNode.left;
+            }
+        }
+        return list;
+    }
+
+    private TreeNode createBST(List<Integer> list, int start, int end) {
+        if (start > end) {
+            return null;
+        }
+        int mid = (start + end) / 2;
+        TreeNode leftTree = createBST(list, start, mid - 1);
+        TreeNode rightTree = createBST(list, mid + 1, end);
+        return new TreeNode(list.get(mid), leftTree, rightTree);
     }
-    inorderTraversal(root.left, sortedList);
-    sortedList.add(root.val);
-    inorderTraversal(root.right, sortedList);
-  }
 }