Problem of The Day: Remove Zero Sum Consecutive Nodes from Linked List

Problem Statement

Intuition

The problem involves removing zero-sum sublists from a linked list. Use stack and hash set data structure to solve the problem.

Approach

I approach the problem by using a stack to maintain the running sum of nodes in the linked list along with the corresponding node. I also use a set (prefix_sum) to keep track of the unique running sums encountered. While iterating through the linked list, if the running sum is already in the set, I remove the subsequence with a zero sum. Finally, I connect the remaining nodes in the stack to form the modified linked list.

Complexity

• Time complexity: O(n), where n is the number of nodes in the linked list. The algorithm iterates through the list once.

• Space complexity: O(n), as the space required for the stack and set is proportional to the number of nodes in the linked list.

Code

# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, val=0, next=None):
#         self.val = val
#         self.next = next
class Solution:
    def removeZeroSumSublists(self, head: Optional[ListNode]) -> Optional[ListNode]:
        dummy = ListNode(-1)
        stack = [[dummy, 0]]
        curr = head
        prefix_sum = {0}
        while curr:
            _, total = stack[-1]
            total += curr.val
            if total in prefix_sum:
                seen = total
                while stack and stack[-1][1] != seen:
                    _, total = stack.pop()
                    prefix_sum.remove(total)
            else:
                stack.append([curr, total])
                prefix_sum.add(total)
            curr = curr.next

        for i in range(len(stack) - 1):
            stack[i][0].next = stack[i + 1][0]

        stack[-1][0].next = None

        return stack[0][0].next

Editorial Solution

Approach 1: Prefix Sum for Each Consecutive Sequence

class Solution:
    def removeZeroSumSublists(self, head: Optional[ListNode]) -> Optional[ListNode]:
        front = ListNode(0, head)
        start = front

        while start is not None:
            prefix_sum = 0
            end = start.next

            while end is not None:
                # Add end's value to the prefix sum
                prefix_sum += end.val
                # Delete zero sum consecutive sequence
                # by setting node before sequence to node after
                if prefix_sum == 0:
                    start.next = end.next
                end = end.next

            start = start.next

        return front.next

• Time complexity: O(n^2)
• Space complexity: O(1)

Approach 2: Prefix Sum Hash Table

implementation visited each node in the linked list twice

class Solution:
    def removeZeroSumSublists(self, head: Optional[ListNode]) -> Optional[ListNode]:
        front = ListNode(0, head)
        current = front
        prefix_sum = 0
        prefix_sum_to_node = {0: front}

        # Calculate the prefix sum for each node and add to the hashmap
        # Duplicate prefix sum values will be replaced
        while current is not None:
            prefix_sum += current.val
            prefix_sum_to_node[prefix_sum] = current
            current = current.next

        # Reset prefix sum and current
        prefix_sum = 0
        current = front

        # Delete zero sum consecutive sequences
        # by setting node before sequence to node after
        while current is not None:
            prefix_sum += current.val
            current.next = prefix_sum_to_node[prefix_sum].next
            current = current.next

        return front.next

Improved implementation

class Solution:
    def removeZeroSumSublists(self, head: Optional[ListNode]) -> Optional[ListNode]:
        front = ListNode(0, head)
        current = front
        prefix_sum = 0
        prefix_sum_to_node = {}
        while current is not None:
            # Add current's value to the prefix sum
            prefix_sum += current.val

            # If prefix_sum is already in the hashmap,
            # we have found a zero-sum sequence:
            if prefix_sum in prefix_sum_to_node:
                prev = prefix_sum_to_node[prefix_sum]
                current = prev.next

                # Delete zero sum nodes from hashmap
                # to prevent incorrect deletions from linked list
                p = prefix_sum + current.val
                while p != prefix_sum:
                    del prefix_sum_to_node[p]
                    current = current.next
                    p += current.val

                # Make connection from the node before
                # the zero sum sequence to the node after
                prev.next = current.next
            else:
                # Add new prefix_sum to hashmap
                prefix_sum_to_node[prefix_sum] = current

            # Progress to next element in list
            current = current.next

        return front.next

• Time complexity: O(n)
• Space complexity: O(n)