Leetcode 1171. Remove Zero Sum Consecutive Nodes from Linked List

Problem Explanation: Given a linked list, the problem asks to delete sequences of consecutive nodes with the total value of 0, and return the head of the final linked list.

We can have multiple answers for the problem. As an instance, if we have a linked list [1,2,-3,3,1] either 3,1 or 1,2,1 could be a valid result. The constraint for the problem is the linked list will have between 1 to 1000 nodes and each node in the linked list has value between -1000 and 1000.

This problem is easier to illustrate using an example. Let's start with the list [1,2,-3,3,1].

At the start: dummy -> 1 -> 2 -> -3 -> 3 -> 1

We store the prefix sum in a HashMap where the key is the sum and the value is the node: {0: dummy, 1: 1, 3: 2, 0: -3, 3: 3, 4: 1}

Then we iterate through the list again, if we see a prefix sum has already been stored in the map, then we know that we have a subarray sums to 0, we just need to delete it: The prefix sum is 0 at -3, and we have stored this prefix sum before at the dummy node, so we remove 1 -> 2 -> -3 and connect dummy to 3. Then the prefix sum is 3 again at 3, but we have stored it at 2 before, so we remove 3 and connect dummy to 1.

Solution:

Python:

1
2python
3class Solution:
4    def removeZeroSumSublists(self, head: ListNode) -> ListNode:
5        # Initialize a dummy node and a container for prefix sum
6        prefixSum = 0
7        prefixSumMap = {0: dummy := ListNode(0)}
8        # Fill up the prefix sum container
9        node = head
10        while node:
11            prefixSum += node.val
12            prefixSumMap[prefixSum] = node
13            node = node.next
14        # Remove zero sum sequences
15        node = dummy
16        prefixSum = 0
17        while node:
18            prefixSum += node.val
19            node.next = prefixSumMap[prefixSum].next
20            node = node.next
21        return dummy.next

Java:

1
2java
3public class Solution {
4    public ListNode removeZeroSumSublists(ListNode head) {
5        HashMap<Integer, ListNode> prefixSumMap = new HashMap<>();
6        ListNode dummy = new ListNode(0);
7        dummy.next = head;
8        prefixSumMap.put(0, dummy);
9        int prefix = 0;
10        while (head != null) {
11            prefix += head.val;
12            if (!prefixSumMap.containsKey(prefix)) {
13                prefixSumMap.put(prefix, head);
14            } else {
15                ListNode node = prefixSumMap.get(prefix).next;
16                int sum = prefix + node.val;
17                while (sum != prefix) {
18                    prefixSumMap.remove(sum);
19                    node = node.next;
20                    sum += node.val;
21                }
22                prefixSumMap.get(prefix).next = head.next;
23            }
24            head = head.next;
25        }
26        return dummy.next;
27    }
28}

JavaScript:

1
2javascript
3var removeZeroSumSublists = function(head) {
4    let prefixSumMap = new Map();
5    let dummy = new ListNode(null);
6    dummy.next = head;
7    prefixSumMap.set(0, dummy);
8    let prefix = 0;
9    
10    while (head != null) {
11        prefix += head.val;
12        if (!prefixSumMap.has(prefix)) {
13            prefixSumMap.set(prefix, head);
14        } else {
15            let node = prefixSumMap.get(prefix).next;
16            let sum = prefix + node.val;
17            while (sum != prefix) {
18                prefixSumMap.delete(sum);
19                node = node.next;
20                if (node) sum += node.val;
21            }
22            prefixSumMap.get(prefix).next = head.next;
23        }
24           head = head.next;
25       }
26    return dummy.next;
27}

C++:

1
2c++
3class Solution {
4public:
5    ListNode* removeZeroSumSublists(ListNode* head) {
6        unordered_map<int, ListNode*> prefixSumMap;
7        ListNode* dummy = new ListNode(0);
8        dummy->next = head;
9        prefixSumMap[0] = dummy;
10        int prefix = 0;
11        while (head != NULL) {
12            prefix += head->val;
13            if (prefixSumMap.count(prefix)) {
14                int sum = prefix;
15                ListNode* node = prefixSumMap[prefix]->next;
16                while (sum != prefix || node != head) {
17                    sum += node->val;
18                    if (prefixSumMap[node->val] == node)
19                        prefixSumMap.erase(sum);
20                    node = node->next;
21                }
22                prefixSumMap[prefix]->next = head->next;
23            } else {
24                prefixSumMap[prefix] = head;
25            }
26            head = head->next;
27        }
28        return dummy->next;
29    }
30};

C#:

1
2csharp
3public class Solution {
4    public ListNode RemoveZeroSumSublists(ListNode head) {
5        Dictionary<int, ListNode> prefixSumMap = new Dictionary<int, ListNode>();
6        ListNode dummy = new ListNode(0, head);
7        prefixSumMap[0] = dummy;
8        int prefix = 0;
9        while (head != null) {
10            prefix += head.val;
11            if (prefixSumMap.ContainsKey(prefix)) {
12                ListNode node = prefixSumMap[prefix].next;
13                int sum = prefix + node.val;
14                while (sum != prefix) {
15                    prefixSumMap.Remove(sum);
16                    node = node.next;
17                    sum += node.val;
18                }
19                prefixSumMap[prefix].next = head.next;
20            } else {
21                prefixSumMap[prefix] = head;
22            }
23            head = head.next;
24        }
25        return dummy.next;
26    }
27}

Rust:

1
2rust
3pub struct ListNode {
4    pub val: i32,
5    pub next: Option<Box<ListNode>>
6}
7
8pub struct Solution;
9
10use std::collections::HashMap;
11impl Solution {
12    pub fn remove_zero_sum_sublists(head: Option<Box<ListNode>>) -> Option<Box<ListNode>> {
13        let mut prefix_sum_map: HashMap<i32, _> = HashMap::new();
14        let mut dummy = Box::new(ListNode{val: 0, next: head});
15        prefix_sum_map.insert(0, &dummy);
16        let mut prefix = 0;
17        let mut node = &mut dummy;
18        while let Some(n) = &mut node.next {
19            prefix += n.val;
20            if prefix_sum_map.contains_key(&prefix) {
21                let mut sum = prefix;
22                let mut to_remove = prefix_sum_map.get(&prefix).unwrap().next.as_ref().map(|n| n.val);
23                while let Some(val) = to_remove {
24                    sum += val;
25                    if prefix_sum_map.get(&sum) != None && sum != prefix {
26                        prefix_sum_map.remove(&sum);
27                    }
28                    to_remove = prefix_sum_map.get(&prefix).unwrap().next.as_ref().map(|n| n.val);
29                }
30                prefix_sum_map.get(&prefix).unwrap().next = n.next.take();
31            } else {
32                prefix_sum_map.insert(prefix, node);
33            }
34            node = n;
35        }
36        dummy.next
37    }
38}

Here, we follow the same approach as above solutions using Rust language. ListNode and Solution are two structures, and we implement remove_zero_sum_sublists function in the Solution structure. First, we initialize the HashMap and ListNode, insert the head of list into HashMap. Then, we traverse the list, keep calculating sum values at each node and check whether the value already exist in HashMap. If it does, then remove these nodes, else add current node into HashMap. Continue this process until reach the end of list.

PHP:

1
2php
3class ListNode {
4    public $val = 0;
5    public $next = null;
6    function __construct($val) { $this->val = $val; }
7}
8
9class Solution {
10    function removeZeroSumSublists($head) {
11        $prefixSumMap = [];
12        $dummy = new ListNode(null);
13        $dummy->next = $head;
14        $prefixSumMap[0] = $dummy;
15        $prefix = 0;
16        
17        while ($head !== null) {
18            $prefix += $head->val;
19            if (!array_key_exists($prefix, $prefixSumMap)) {
20                $prefixSumMap[$prefix] = $head;
21            } else {
22                $node = $prefixSumMap[$prefix]->next;
23                $sum = $prefix + $node->val;
24                while ($sum !== $prefix) {
25                    unset($prefixSumMap[$sum]);
26                    $node = $node->next;
27                    if ($node) $sum += $node->val;
28                }
29                $prefixSumMap[$prefix]->next = $head->next;
30            }
31
32            $head = $head->next;
33        }
34        return $dummy->next;
35    }
36}

Similar logic is applied in the PHP language. We have the ListNode class for linked list nodes, and we implement the removeZeroSumSublists function in the Solution class. We use an associative array to store a node with sum value as a key. If the sum value finds a existing one in the associative array, we remove those nodes in the linked list. We continue this operation for all nodes in the linked list, and return the result.