Modified 2 solutions in trees

Author: varunu28

Trees/Problems/Inorder Traversal.java

@@ -12,27 +12,22 @@
  * }
  */
 public class Solution {
-    public ArrayList<Integer> inorderTraversal(TreeNode A) {
-        ArrayList<Integer> arr = new ArrayList<>();
-        Stack<TreeNode> stack = new Stack<>();
-       
-        if (A == null) return arr;
-        
-        while (true) {
-            if (A != null) {
-            	stack.push(A);
-            	A = A.left;
-            }
-            else {
-            	if (stack.isEmpty()) {
-            		break;
-            	}
-            	A = stack.pop();
-            	arr.add(A.val);
-            	A = A.right;
-            }
-        }
-        
-        return arr;
+  public ArrayList<Integer> inorderTraversal(TreeNode A) {
+    ArrayList<Integer> list = new ArrayList<>();
+    Stack<TreeNode> stack = new Stack<>();
+    while (A != null) {
+      stack.push(A);
+      A = A.left;
     }
+    while (!stack.isEmpty()) {
+      TreeNode removed = stack.pop();
+      list.add(removed.val);
+      TreeNode right = removed.right;
+      while (right != null) {
+        stack.push(right);
+        right = right.left;
+      }
+    }
+    return list;
+  }
 }

Trees/Problems/Invert The Binary Tree.java

@@ -12,21 +12,15 @@
  * }
  */
 public class Solution {
-    public TreeNode invertTree(TreeNode A) {
-        if (A == null) return A;
-        swap(A);
-        return A;
-    }
-    
-    public void swap(TreeNode root) {
-        if (root.left == null && root.right == null) {
-            return;
-        }
-        TreeNode n1 = root.left;
-        root.left = root.right;
-        root.right = n1;
-        
-        if (root.left != null) swap(root.left);
-        if (root.right != null) swap(root.right);
+  public TreeNode invertTree(TreeNode A) {
+    if (A == null) {
+      return null;
     }
+    TreeNode left = A.left;
+    A.left = A.right;
+    A.right = left;
+    invertTree(A.left);
+    invertTree(A.right);
+    return A;
+  }
 }