edit 538

fishercoder1534 · fishercoder1534 · commit 9f71db46940d · 2017-08-09T08:43:44.000-07:00
diff --git a/README.md b/README.md
@@ -101,7 +101,7 @@ Your ideas/fixes/algorithms are more than welcome!
 |541|[Reverse String II](https://leetcode.com/problems/reverse-string-ii/)|[Solution](../master/src/main/java/com/fishercoder/solutions/_541.java) | O(n) |O(1) | Easy | String
 |540|[Single Element in a Sorted Array](https://leetcode.com/problems/single-element-in-a-sorted-array/)|[Solution](../master/src/main/java/com/fishercoder/solutions/_540.java) | O(n) |O(1) | Medium | 
 |539|[Minimum Time Difference](https://leetcode.com/problems/minimum-time-difference/)|[Solution](../master/src/main/java/com/fishercoder/solutions/MinimumTimeDifference.java) | O(n) |O(1) | Medium | String
-|538|[Convert BST to Greater Tree](https://leetcode.com/problems/convert-bst-to-greater-tree/)|[Solution](../master/src/main/java/com/fishercoder/solutions/_538.java) | O(nlogn) |O(n) | Medium | Tree
+|538|[Convert BST to Greater Tree](https://leetcode.com/problems/convert-bst-to-greater-tree/)|[Solution](../master/src/main/java/com/fishercoder/solutions/_538.java) | O(n) |O(h) | Medium | Tree
 |537|[Complex Number Multiplication](https://leetcode.com/problems/complex-number-multiplication/)|[Solution](../master/src/main/java/com/fishercoder/solutions/_537.java) | O(1) |O(1) | Medium | Math, String
 |536|[Construct Binary Tree from String](https://leetcode.com/problems/construct-binary-tree-from-string/)|[Solution](../master/src/main/java/com/fishercoder/solutions/_536.java) | O(n) |O(h) | Medium | Recursion
 |535|[Encode and Decode TinyURL](https://leetcode.com/problems/encode-and-decode-tinyurl/)|[Solution](../master/src/main/java/com/fishercoder/solutions/_535.java) | O(1) |O(n) | Medium | Design
diff --git a/src/main/java/com/fishercoder/solutions/_538.java b/src/main/java/com/fishercoder/solutions/_538.java
@@ -27,42 +27,62 @@
  */
 
 public class _538 {
-    //This solution is generic for both BST and regular binary trees
-    public TreeNode convertBST(TreeNode root) {
-        if (root == null) return root;
-        List<Integer> list = new ArrayList<>();
-        putNodeToList(list, root);
-        Collections.sort(list);
-        int[] sums = new int[list.size()];
-        sums[list.size()-1] = 0;
-        for (int i = list.size()-2; i >= 0; i--) {
-            sums[i] = sums[i+1] + list.get(i+1);
+
+    public static class BSTSolution {
+        /**reference: https://discuss.leetcode.com/topic/83455/java-recursive-o-n-time/13
+         *
+         * A reversed version of inorder traversal: right -> root -> left*/
+        public TreeNode convertBST(TreeNode root) {
+            dfs(root, 0);
+            return root;
         }
-        TreeMap<Integer, Integer> treeMap = new TreeMap<>();
-        for (int i = 0; i < list.size(); i++) {
-            treeMap.put(list.get(i), sums[i]);
+
+        private int dfs(TreeNode root, int val) {
+            if (root == null) return val;
+            int right = dfs(root.right, val);
+            root.val += right;
+            int left = dfs(root.left, root.val);
+            return left;
         }
-        TreeNode result = new TreeNode(treeMap.get(list.get(0)));
-        return generateResultRoot(root, treeMap, result);
     }
 
-    private TreeNode generateResultRoot(TreeNode root, TreeMap<Integer, Integer> treeMap, TreeNode result) {
-        if (root != null) result.val = treeMap.get(root.val) + root.val;
-        if (root.left != null) {
-            result.left = new TreeNode(0);
-            generateResultRoot(root.left, treeMap, result.left);
+    public static class GenericSolution {
+        //This solution is generic for both BST and regular binary trees
+        public TreeNode convertBST(TreeNode root) {
+            if (root == null) return root;
+            List<Integer> list = new ArrayList<>();
+            putNodeToList(list, root);
+            Collections.sort(list);
+            int[] sums = new int[list.size()];
+            sums[list.size() - 1] = 0;
+            for (int i = list.size() - 2; i >= 0; i--) {
+                sums[i] = sums[i + 1] + list.get(i + 1);
+            }
+            TreeMap<Integer, Integer> treeMap = new TreeMap<>();
+            for (int i = 0; i < list.size(); i++) {
+                treeMap.put(list.get(i), sums[i]);
+            }
+            TreeNode result = new TreeNode(treeMap.get(list.get(0)));
+            return generateResultRoot(root, treeMap, result);
         }
-        if (root.right != null) {
-            result.right = new TreeNode(0);
-            generateResultRoot(root.right, treeMap, result.right);
+
+        private TreeNode generateResultRoot(TreeNode root, TreeMap<Integer, Integer> treeMap, TreeNode result) {
+            if (root != null) result.val = treeMap.get(root.val) + root.val;
+            if (root.left != null) {
+                result.left = new TreeNode(0);
+                generateResultRoot(root.left, treeMap, result.left);
+            }
+            if (root.right != null) {
+                result.right = new TreeNode(0);
+                generateResultRoot(root.right, treeMap, result.right);
+            }
+            return result;
         }
-        return result;
-    }
 
-    private void putNodeToList(List<Integer> list, TreeNode root) {
-        if (root != null) list.add(root.val);
-        if (root.left != null) putNodeToList(list, root.left);
-        if (root.right != null) putNodeToList(list, root.right);
+        private void putNodeToList(List<Integer> list, TreeNode root) {
+            if (root != null) list.add(root.val);
+            if (root.left != null) putNodeToList(list, root.left);
+            if (root.right != null) putNodeToList(list, root.right);
+        }
     }
-
 }
diff --git a/src/test/java/com/fishercoder/_538Test.java b/src/test/java/com/fishercoder/_538Test.java
@@ -9,14 +9,16 @@
 import static junit.framework.Assert.assertEquals;
 
 public class _538Test {
-    private static _538 test;
+    private static _538.GenericSolution genericSolution;
+    private static _538.BSTSolution bstSolution;
     private static TreeNode actualRoot;
     private static TreeNode expectedRoot;
     private static TreeNode root;
 
     @BeforeClass
     public static void setup(){
-        test = new _538();
+        bstSolution = new _538.BSTSolution();
+        genericSolution = new _538.GenericSolution();
     }
 
     @Before
@@ -28,20 +30,18 @@ public void test1(){
         root = new TreeNode(5);
         root.left = new TreeNode(2);
         root.right = new TreeNode(13);
-
         expectedRoot = new TreeNode(18);
         expectedRoot.left = new TreeNode(20);
         expectedRoot.right = new TreeNode(13);
-        actualRoot = test.convertBST(root);
-        assertEquals(expectedRoot.toString(), actualRoot.toString());
+        assertEquals(expectedRoot.toString(), genericSolution.convertBST(root).toString());
+        assertEquals(expectedRoot.toString(), bstSolution.convertBST(root).toString());
     }
 
     @Test
     public void test2(){
         root = null;
-
         expectedRoot = null;
-        actualRoot = test.convertBST(root);
-        assertEquals(expectedRoot, actualRoot);
+        assertEquals(expectedRoot, genericSolution.convertBST(root));
+        assertEquals(expectedRoot, bstSolution.convertBST(root));
     }
 }
