Remove Zero Sum Consecutive Nodes from Linked List - LeetCode Daily Challenge

Problem Statement

Given the head of a linked list, we repeatedly delete consecutive sequences of nodes that sum to 0 until there are no such sequences.

After doing so, return the head of the final linked list.  You may return any such answer.

 (Note that in the examples below, all sequences are serializations of ListNode objects.)

Example 1

Input: head = [1,2,-3,3,1]
Output: [3,1]
Note: The answer [1,2,1] would also be accepted.

Example 2

Input: head = [1,2,3,-3,4]
Output: [1,2,4]

Try here before watching the video.

Video Solution

Java Code

class Solution {
    public ListNode removeZeroSumSublists(ListNode head) {
        Map<Integer, ListNode> prefixSumMap = new HashMap<>();
        ListNode dummy = new ListNode(0);
        dummy.next = head;
        prefixSumMap.put(0, dummy);
        ListNode ptr = head;
        int prefixSum = 0;
        
        while(ptr != null) {
            prefixSum += ptr.val;
            if(prefixSumMap.containsKey(prefixSum)) {
                ListNode prev = prefixSumMap.get(prefixSum);
                ListNode temp = prev.next;
                int sum = prefixSum;
                while(temp != ptr) {
                    sum += temp.val;
                    prefixSumMap.remove(sum);
                    temp = temp.next;
                }
                prev.next = ptr.next;
            } else {
                prefixSumMap.put(prefixSum, ptr);
            }
            ptr = ptr.next;
        }
        return dummy.next;
    }
}

C++ Code

class Solution {
public:
    ListNode* removeZeroSumSublists(ListNode* head) {
        unordered_map<int, ListNode*> prefixSumMap;
        ListNode* dummy = new ListNode(0);
        dummy->next = head;
        prefixSumMap[0] = dummy;
        ListNode* ptr = head;
        int prefixSum = 0;

        while (ptr != nullptr) {
            prefixSum += ptr->val;
            if (prefixSumMap.find(prefixSum) != prefixSumMap.end()) {
                ListNode* prev = prefixSumMap[prefixSum];
                ListNode* temp = prev->next;
                int sum = prefixSum;
                while (temp != ptr) {
                    sum += temp->val;
                    prefixSumMap.erase(sum);
                    temp = temp->next;
                }
                prev->next = ptr->next;
            } else {
                prefixSumMap[prefixSum] = ptr;
            }
            ptr = ptr->next;
        }
        return dummy->next;
    }
};

Python Code

class Solution:
    def removeZeroSumSublists(self, head: ListNode) -> ListNode:
        prefix_sum_map = dict()
        dummy = ListNode(0)
        dummy.next = head
        ptr = head
        prefix_sum = 0
        prefix_sum_map[0] = dummy

        while ptr:
            prefix_sum += ptr.val
            if prefix_sum in prefix_sum_map:
                prev = prefix_sum_map[prefix_sum]
                temp = prev.next
                _sum = prefix_sum
                while temp != ptr:
                    _sum += temp.val
                    del prefix_sum_map[_sum]
                    temp = temp.next
                prev.next = ptr.next
            else:
                prefix_sum_map[prefix_sum] = ptr
            ptr = ptr.next

        return dummy.next

Javascript Code

var removeZeroSumSublists = function(head) {
    let prefixSumMap = new Map();
    let dummy = new ListNode(0);
    dummy.next = head;
    prefixSumMap.set(0, dummy);
    let ptr = head;
    let prefixSum = 0;

    while (ptr !== null) {
        prefixSum += ptr.val;
        if (prefixSumMap.has(prefixSum)) {
            let prev = prefixSumMap.get(prefixSum);
            let temp = prev.next;
            let sum = prefixSum;
            while (temp !== ptr) {
                sum += temp.val;
                prefixSumMap.delete(sum);
                temp = temp.next;
            }
            prev.next = ptr.next;
        } else {
            prefixSumMap.set(prefixSum, ptr);
        }
        ptr = ptr.next;
    }
    return dummy.next;
};

Go Code

func removeZeroSumSublists(head *ListNode) *ListNode {
    prefixSumMap := make(map[int]*ListNode)
    dummy := &ListNode{Val: 0, Next: head}
    prefixSumMap[0] = dummy
    ptr := head
    prefixSum := 0

    for ptr != nil {
        prefixSum += ptr.Val
        if prev, ok := prefixSumMap[prefixSum]; ok {
            temp := prev.Next
            sum := prefixSum
            for temp != ptr {
                sum += temp.Val
                delete(prefixSumMap, sum)
                temp = temp.Next
            }
            prev.Next = ptr.Next
        } else {
            prefixSumMap[prefixSum] = ptr
        }
        ptr = ptr.Next
    }
    return dummy.Next
}

Complexity Analysis

Time Complexity: O(N), we are iterating through each node.
Space Complexity: O(N), in the worst-case scenario, all the numbers in the linked list might be positive and we may end up storing the N number of different sums in prefixSumMap.