refactor 226

fishercoder1534 · fishercoder1534 · commit 72c44f247627 · 2021-09-23T10:27:38.000-07:00
diff --git a/src/main/java/com/fishercoder/solutions/_226.java b/src/main/java/com/fishercoder/solutions/_226.java
@@ -8,8 +8,11 @@
 public class _226 {
 
     public static class Solution1 {
+        /**
+         * An iterative solution
+         */
         public TreeNode invertTree(TreeNode root) {
-            if (root == null) {
+            if (root == null || (root.left == null && root.right == null)) {
                 return root;
             }
             Queue<TreeNode> q = new LinkedList();
@@ -31,8 +34,11 @@ public TreeNode invertTree(TreeNode root) {
     }
 
     public static class Solution2 {
+        /**
+         * A recursive solution
+         */
         public TreeNode invertTree(TreeNode root) {
-            if (root == null) {
+            if (root == null || (root.left == null && root.right == null)) {
                 return root;
             }
             TreeNode temp = root.left;
@@ -45,9 +51,11 @@ public TreeNode invertTree(TreeNode root) {
     }
 
     public static class Solution3 {
-        //more concise version
+        /**
+         * A more concise version
+         */
         public TreeNode invertTree(TreeNode root) {
-            if (root == null) {
+            if (root == null || (root.left == null && root.right == null)) {
                 return root;
             }
             TreeNode temp = root.left;
