【C++】連結リストが二分木の下向きパスと一致するかを判定するアルゴリズム

二分木のルート(root)と、先頭ノードheadを持つ連結リストが与えられたとします。連結リストのhead以降のすべての要素が、二分木内のどこかの下向きパス(downward path)に一致する場合はTrueを、一致しない場合はFalseを返す必要があります。

例えば、次のような二分木があったとします。

このとき、連結リストが [1, 4, 2, 6] であれば、出力は true になります。実際に、ルートの1から始まり4→2→6とたどるパスが存在するためです。

解法のアプローチ

この問題を解くために、再帰とメモ化(動的計画法)を組み合わせた以下の手順に従います。

• メモ化用のマップ dp を定義します。
• head・root・flag の3つの引数を受け取る solve() メソッドを定義します。
• head が NULL の場合は true を返し、root が NULL の場合は false を返します。
• dp に head が存在し、かつ dp[head] に root が存在し、さらに dp[head][root] に flag が存在する場合は、計算済みの dp[head][root][flag] をそのまま返します(メモ化による高速化)。
• head の値と root の値が等しい場合:
  • ret := solve(head の next, root の左部分木, false) OR solve(head の next, root の右部分木, false) を計算します。
  • ret が真であれば、dp[head][root][flag] := true として返します。
  • そうでなければ、dp[head][root][flag] := solve(head, root の左部分木, flag) OR solve(head, root の右部分木, flag) とします。
  • 最後に dp[head][root][flag] を返します。
• 値が異なり、かつ flag が立っていない場合は、dp[head][root][flag] := false を返します。
• それ以外の場合は、dp[head][root][flag] := solve(head, root の左部分木, flag) OR solve(head, root の右部分木, flag) を返します。
• main 関数からは solve(head, root, true) を呼び出します。

ここで flag は「現在、連結リストとのマッチング途中である」ことを表しています。マッチングの途中で値が不一致になった場合は即座に false を返しますが、まだマッチングを開始していない状態であれば、子ノード側で新たなマッチング開始点を探し続けることができます。

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;
}
class TreeNode{
    public:
        int val;
        TreeNode *left, *right;
        TreeNode(int data){
            val = data;
            left = right = NULL;
        }
};
void insert(TreeNode **root, int val){
    queue<TreeNode*> q;
    q.push(*root);
    while(q.size()){
        TreeNode *temp = q.front();
        q.pop();
        if(!temp->left){
            if(val != NULL)
                temp->left = new TreeNode(val);
            else
                temp->left = new TreeNode(0);
            return;
        } else {
            q.push(temp->left);
        }
        if(!temp->right){
            if(val != NULL)
                temp->right = new TreeNode(val);
            else
                temp->right = new TreeNode(0);
            return;
        } else {
            q.push(temp->right);
        }
    }
}
TreeNode *make_tree(vector<int> v){
    TreeNode *root = new TreeNode(v[0]);
    for(int i = 1; i<v.size(); i++){
        insert(&root, v[i]);
    }
    return root;
}
class Solution {
    public:
        map < ListNode*, map<TreeNode*, map <bool, bool>> > dp;
        bool solve(ListNode* head, TreeNode* root, bool flag = true){
            if(head == NULL) return true;
            if(!root) return false;
            if(dp.count(head) && dp[head].count(root) && dp[head]
               [root].count(flag)) return dp[head][root][flag];
            if(head->val == root->val){
                bool ret = solve(head->next, root->left, false) ||
                solve(head->next, root->right, false);
                if(ret) return dp[head][root][flag] = true;
                return dp[head][root][flag] = solve(head, root->left,
                 flag) || solve(head, root->right, flag);
            }else if(!flag) return dp[head][root][flag] = false;
            else
                return dp[head][root][flag]= solve(head, root->left,
                 flag) || solve(head, root->right, flag);
        }
        bool isSubPath(ListNode* head, TreeNode* root) {
            return solve(head, root);
        }
};
main(){
    vector<int> v = {1,4,2,6};
    vector<int> v1 = {1,4,4,NULL,2,2,NULL,1,NULL,6,8,NULL,NULL,NULL,NULL,1,3};
    ListNode *head = make_list(v);
    TreeNode *root = make_tree(v1);
    Solution ob;
    cout << (ob.isSubPath(head, root));
}

入力

[1,4,2,6]
[1,4,4,null,2,2,null,1,null,6,8,null,null,null,null,1,3]

出力

1

出力が 1(true)となり、連結リスト [1,4,2,6] が二分木内の下向きパスとして存在することが確認できます。この手法では、各ノードをマッチング開始点として試しつつ、メモ化によって同じ状態(head, root, flag の組み合わせ)の再計算を避けているため、効率的に答えを求めることができます。