add a second solution for 1171

fishercoder1534 · fishercoder1534 · commit 2bc223748252 · 2021-02-02T17:03:58.000-08:00
diff --git a/src/main/java/com/fishercoder/solutions/_1171.java b/src/main/java/com/fishercoder/solutions/_1171.java
@@ -72,4 +72,37 @@ private List<Integer> shrinkList(List<Integer> list) {
             return list;
         }
     }
+
+    public static class Solution2 {
+        /**
+         * credit: https://leetcode.com/problems/remove-zero-sum-consecutive-nodes-from-linked-list/discuss/366337/Java-Iterative-and-Recursive-solution
+         * this post explains it all
+         * key of the hashmap is the prefix sum of all the nodes we've gone so far
+         * value of the hashmap is the corresponding linked list node
+         */
+        public ListNode removeZeroSumSublists(ListNode head) {
+            ListNode pre = new ListNode(-1);
+            ListNode curr = pre;
+            pre.next = head;
+            Map<Integer, ListNode> map = new HashMap<>();
+            int preSum = 0;
+            while (curr != null) {
+                preSum += curr.val;
+                if (map.containsKey(preSum)) {
+                    curr = map.get(preSum).next;
+                    int key = preSum + curr.val;
+                    while (key != preSum) {
+                        map.remove(key);
+                        curr = curr.next;
+                        key += curr.val;
+                    }
+                    map.get(preSum).next = curr.next;
+                } else {
+                    map.put(preSum, curr);
+                }
+                curr = curr.next;
+            }
+            return pre.next;
+        }
+    }
 }
diff --git a/src/test/java/com/fishercoder/_1171Test.java b/src/test/java/com/fishercoder/_1171Test.java
@@ -10,20 +10,24 @@
 
 public class _1171Test {
     private static _1171.Solution1 solution1;
+    private static _1171.Solution2 solution2;
 
     @BeforeClass
     public static void setup() {
         solution1 = new _1171.Solution1();
+        solution2 = new _1171.Solution2();
     }
 
     @Test
     public void test1() {
         assertEquals(LinkedListUtils.contructLinkedList(new int[]{3, 1}), solution1.removeZeroSumSublists(LinkedListUtils.contructLinkedList(new int[]{1, 2, -3, 3, 1})));
+        assertEquals(LinkedListUtils.contructLinkedList(new int[]{3, 1}), solution2.removeZeroSumSublists(LinkedListUtils.contructLinkedList(new int[]{1, 2, -3, 3, 1})));
     }
 
     @Test
     public void test2() {
         assertEquals(LinkedListUtils.contructLinkedList(new int[]{1, 2, 4}), solution1.removeZeroSumSublists(LinkedListUtils.contructLinkedList(new int[]{1, 2, 3, -3, 4})));
+        assertEquals(LinkedListUtils.contructLinkedList(new int[]{1, 2, 4}), solution2.removeZeroSumSublists(LinkedListUtils.contructLinkedList(new int[]{1, 2, 3, -3, 4})));
     }
 
     @Test
