refactor 666

Author: fishercoder1534

src/main/java/com/fishercoder/solutions/_666.java

@@ -2,6 +2,9 @@
 
 import com.fishercoder.common.classes.TreeNode;
 
+import java.util.HashMap;
+import java.util.Map;
+
 /**
  * 666. Path Sum IV
  * If the depth of a tree is smaller than 5, then this tree can be represented by a list of three-digits integers.
@@ -16,7 +19,7 @@
  The units digit represents the value V of this node, 0 <= V <= 9.
 
  Given a list of ascending three-digits integers representing a binary with the depth smaller than 5.
- You need to return the sum of all paths from the root towards the leaves.
+ You need to return the totalSum of all paths from the root towards the leaves.
 
  Example 1:
 
@@ -28,7 +31,7 @@
   / \
  5   1
 
- The path sum is (3 + 5) + (3 + 1) = 12.
+ The path totalSum is (3 + 5) + (3 + 1) = 12.
 
  Example 2:
 
@@ -40,7 +43,7 @@ The path sum is (3 + 5) + (3 + 1) = 12.
   \
   1
 
- The path sum is (3 + 1) = 4.
+ The path totalSum is (3 + 1) = 4.
 
  */
 public class _666 {
@@ -125,4 +128,38 @@ private TreeNode constructTree(int[] nums) {
             return root;
         }
     }
+
+    public static class Solution2 {
+        public int totalSum = 0;
+
+        public int pathSum(int[] nums) {
+            Map<Integer, Integer> map = new HashMap<>();
+            for (int i = 0; i < nums.length; i++) {
+                int key = nums[i] / 10;
+                int value = nums[i] % 10;
+                map.put(key, value);
+            }
+            dfs(nums[0] / 10, 0, map);
+            return totalSum;
+        }
+
+        private void dfs(int node, int preSum, Map<Integer, Integer> map) {
+            int level = node / 10;
+            int pos = node % 10;
+            int leftChild = (level + 1) * 10 + pos * 2 - 1;
+            int rightChild = (level + 1) * 10 + pos * 2;
+            int currSum = preSum + map.get(node);
+            if (!map.containsKey(leftChild) && !map.containsKey(rightChild)) {
+                totalSum += currSum;
+                return;
+            }
+
+            if (map.containsKey(leftChild)) {
+                dfs(leftChild, currSum, map);
+            }
+            if (map.containsKey(rightChild)) {
+                dfs(rightChild, currSum, map);
+            }
+        }
+    }
 }

src/test/java/com/fishercoder/_666Test.java

@@ -1,39 +1,48 @@
 package com.fishercoder;
 
 import com.fishercoder.solutions._666;
+import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
 
 import static org.junit.Assert.assertEquals;
 
 public class _666Test {
     private static _666.Solution1 solution1;
+    private static _666.Solution2 solution2;
     private static int[] nums;
 
     @BeforeClass
     public static void setup() {
         solution1 = new _666.Solution1();
+        solution2 = new _666.Solution2();
+    }
+
+    @Before
+    public void cleanUp() {
+        solution1.totalSum = 0;
+        solution2.totalSum = 0;
     }
 
     @Test
     public void test1() {
         nums = new int[]{113, 215, 221};
-        solution1.totalSum = 0;
         assertEquals(12, solution1.pathSum(nums));
+        assertEquals(12, solution2.pathSum(nums));
     }
 
     @Test
     public void test2() {
         nums = new int[]{113, 221};
-        solution1.totalSum = 0;
         assertEquals(4, solution1.pathSum(nums));
+        assertEquals(4, solution2.pathSum(nums));
     }
 
     @Test
     public void test3() {
         nums = new int[]{113, 214, 221, 348, 487};
-        solution1.totalSum = 0;
         assertEquals(26, solution1.pathSum(nums));
+        assertEquals(26, solution2.pathSum(nums));
     }
 
 }