C++で連結リストから合計が0になる連続ノードを削除する方法

連結リストの先頭ノード（head）が与えられたとき、ノード値の合計が0になる連続したノード列を、そのような列が一切なくなるまで繰り返し削除し、最終的な連結リストの先頭を返すことを考えます。
例えば、連結リストが [1, 2, -3, 3, 1] という構成の場合、「1 + 2 + (-3) = 0」となるため先頭の3ノードが削除され、答えは [3, 1] になります。

解決の手順

この問題は、プレフィックス和（累積和）とハッシュマップを組み合わせることで効率的に解くことができます。アルゴリズムの流れは以下の通りです。

• 値0を持つダミーノード（dummy）を作成し、dummy の next を head に接続します。

• マップ m を作成し、キー0に dummy を登録します。累積和 sum を0で初期化します。

• head がNULLでない間、次の処理を繰り返します。

  • sum に head の値を加算し、m[sum] に head を登録してから、head を次のノードへ進めます。

• head を dummy に戻し、sum を0にリセットします。

• 再び head がNULLでない間、以下の処理を行います。

  • sum に head の値を加算します。

  • temp に m[sum] を取得します。

  • temp が head 自身と異なる場合、head の next を temp の next へつなぎ替えます。これにより、合計が0になる区間が一括で削除されます。

  • head を次のノードへ進めます。

• 最後に dummy の next を返します。

このアルゴリズムが機能する理由

ポイントはプレフィックス和の性質です。リストを先頭から走査した際、同じプレフィックス和が2つの位置で現れた場合、その間にあるノードの値の総和は必ず0になります。1回目の走査では、各プレフィックス和に対応する最後尾の出現ノードをマップに記録しておきます。2回目の走査では、同じプレフィックス和が見つかった時点でそのノードへ直接つなぎ替えることで、間のゼロサム区間をまとめて取り除いています。

リスト長を n とすると、時間計算量・空間計算量はいずれも O(n) です。

C++による実装例

理解を深めるために、以下の実装を確認してみましょう。

#include <bits/stdc++.h>
using namespace std;
class ListNode{
    public:
    int val;
    ListNode *next;
    ListNode(int data){
        val = data;
        next = NULL;
    }
};
ListNode *make_list(vector<int> v){
    ListNode *head = new ListNode(v[0]);
    for(int i = 1; i<v.size(); i++){
        ListNode *ptr = head;
        while(ptr->next != NULL){
            ptr = ptr->next;
        }
        ptr->next = new ListNode(v[i]);
    }
    return head;
}
void print_list(ListNode *head){
    ListNode *ptr = head;
    cout << "[";
    while(ptr){
        cout << ptr->val << ", ";
        ptr = ptr->next;
    }
    cout << "]" << endl;
}
class Solution {
    public:
    ListNode* removeZeroSumSublists(ListNode* head) {
        ListNode* dummy = new ListNode(0);
        dummy->next = head;
        unordered_map <int, ListNode*> m;
        m[0] = dummy;
        int sum = 0;
        while(head){
            sum += head->val;
            m[sum] = head;
            head = head->next;
        }
        head = dummy;
        sum = 0;
        while(head){
            sum += head->val;
            ListNode* temp = m[sum];
            if(temp != head){
                head->next = temp->next;
            }
            head = head->next;
        }
        return dummy->next;
    }
};
main(){
    vector<int> v1 = {1,2,-3,3,1};
    ListNode *head = make_list(v1);
    Solution ob;
    print_list(ob.removeZeroSumSublists(head));
}

入力

[1,2,-3,3,1]

出力

[3,1]