[LeetCode] 1171. Remove Zero Sum Consecutive Nodes from Linked List

https://leetcode.com/problems/remove-zero-sum-consecutive-nodes-from-linked-list/description/

Algorithm

와...진짜 끙끙댔던 문제이다. 개인적으로 아이디어 떠올리기가 힘들었다. 결국 베스트 풀이를 봤는데...해시맵과 prefixsum의 결합이라니! 상상도 못했음...갈길이 참 멀군...

알고리즘은 여기를 참고했다. 방법은 다음과같다. 노드를 순회하며 prefix_sum을 해시맵에 저장해나간다.

이때

map에 prefixSum이 없다면 map에 저장하고,

이미 있다면 중간의 노드들을 제거하고, 중복된첫노드와 중복된두번째노드의 다음노드를 잇는다.(위 연두색 참고)

같은 prefixSum이 나왔다는건, 결국 중간의 합이 0이라는 뜻이다. 

이때, 엣지케이스를 유의해야한다. 예를들어 1,3,-4가 있을때 모두를 합하면 0이된다.초기 모든 노드합이 0인경우를 고려하기 위해 맨앞에 더미노드 0을 추가한다.

좀더 상세히 수도코드를 적어봤다.

위 그림처럼 head를 이동하다가 같은 prefixSum이 나오면, start를 중복된첫번째노드로 설정하고 temp,psum을 통해 해시맵에서 중간노드들을 제거해나간다. 이 중간노드들을 제거하는 이유는 예를들어 위에서 첫번째 5탐색시, 이전에 삭제한 노드범위를 중복해서 탐지하지 않게 하기 위해서다.

복잡해보이지만 차근차근 곱씹고 이해하니 알겠다.

---

Python

# 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]:
        prefixSum = 0 # 현재까지 누적합
        mp = {} # [누적합 : 노드] 저장할 해시맵
        
        # 초기 누적 합 0에 대응하는 더미 노드
        dummy = ListNode(0)
        dummy.next = head
        mp[0] = dummy

        # 연결리스트 순회
        while head:
            prefixSum += head.val

            # 현재 누적합이 이미 해시 맵에 존재하는 경우
            if prefixSum in mp:
                # 중복된 누적합의 이전 노드를 찾음
                start = mp[prefixSum]
                temp = start
                pSum = prefixSum

                # 중복된 누적합까지의 노드를 해시 맵에서 제거하며 순회
                while temp!= head:
                    temp = temp.next # 다음 노드로 이동
                    pSum += temp.val # 현재노드의 값을 누적합에 더함
                    if temp != head:
                        del mp[pSum] # 중복된 누적 합 해시맵에서 제거
                start.next = head.next
            else:
                # 해시 맵에 현재 누적합 , 노드를 저장
                mp[prefixSum] = head
            
            # 다음 노드로 이동
            head = head.next
        # 더미 노드 다음부터 시작하는 최종 결과 반환
        return dummy.next

Java

/**
 * Definition for singly-linked list.
 * public class ListNode {
 *     int val;
 *     ListNode next;
 *     ListNode() {}
 *     ListNode(int val) { this.val = val; }
 *     ListNode(int val, ListNode next) { this.val = val; this.next = next; }
 * }
 */
class Solution {
    public ListNode removeZeroSumSublists(ListNode head) {
        int prefixSum = 0;
        HashMap<Integer, ListNode> map = new HashMap<>();
        ListNode dummy = new ListNode(0);
        dummy.next = head;
        map.put(0,dummy); // {prefixSum : Node}
        while (head != null){
            prefixSum += head.val;
            if(map.containsKey(prefixSum)){ // 만약 prefixSum이 해시맵에 이미 있다면
                ListNode start = map.get(prefixSum);
                ListNode temp = start;
                int sum = prefixSum;
                while(temp != head){ // 순회하며 노드 맵에서 제거
                    temp = temp.next;
                    sum += temp.val;
                    if(temp != head){map.remove(sum);}
                }
                start.next = head.next;
            } else {
                map.put(prefixSum,head);
            }
            head = head.next; // 순회(다음 노드로)
        }
        return dummy.next;
    }
}