C++で二分木のターゲットノードから距離Kにあるすべてのノードを求める方法

問題概要

二分木・ターゲットノード・値Kが与えられたとき、ターゲットノードから距離Kにあるすべてのノードの値をリストとして返す問題を考えます。

たとえば、入力が root = [3,5,1,6,2,0,8,null,null,7,4]、target = 5、K = 2 の場合、出力は [7,4,1] となります。これは、ターゲットノードから距離2の位置にあるノードの値が 7、4、1 だからです。

解法の考え方

この問題を解くには、まずDFS(深さ優先探索)で各ノードの親へのポインタを記録し、その後ターゲットノードを起点としてBFS(幅優先探索)を行います。二分木は通常、子から親へ辿れないため、事前に親情報を用意しておくのがポイントです。

ステップ1: DFSで親ノードを記録する

• dfs(node, pa) 関数を定義します(pa の初期値は NULL)。
• node が NULL の場合は何もせずに return します。
• parent[node] に pa を代入し、親ノードを記録します。
• 左の子に対して dfs(node->left, node) を再帰的に呼び出します。
• 右の子に対して dfs(node->right, node) を再帰的に呼び出します。

ステップ2: BFSで距離Kのノードを収集する

メイン処理では、以下の手順で探索を進めます。

• 結果を格納する配列 ans を定義します。
• dfs(root) を呼び出して親情報を構築します。
• (ノード, 深度) のペアを保持するキュー q を用意します。
• { target, 0 } を q に挿入します。
• 訪問済みノードを管理する集合 visited を定義し、target を追加します。
• q が空になるまで、以下を繰り返します。
  • q の先頭要素を取り出し、深度 level とノード node を取得します。
  • level == k であれば、node の値を ans の末尾に追加します。
  • 左の子が存在し、level + 1 <= k かつ未訪問であれば、{ 左の子, level + 1 } を q に挿入し、visited にも追加します。
  • 右の子についても同様に処理します。
  • parent[node] が存在し、level + 1 <= k かつ未訪問であれば、{ parent[node], level + 1 } を q に挿入し、visited に追加します。
• 最後に ans を返します。

探索時には「左の子・右の子・親」の3方向へ展開できる点が特徴です。visited 集合によって同じノードを二度訪問するのを防ぎ、無限ループや重複を回避しています。

C++による実装例

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

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<int> v){
    cout << "[";
    for(int i = 0; i<v.size(); i++){
        cout << v[i] << ", ";
    }
    cout << "]"<<endl;
}
class TreeNode{
public:
    int val;
    TreeNode *left, *right;
    TreeNode(int data){
        val = data;
        left = NULL;
        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 <TreeNode*, TreeNode*> parent;
    void dfs(TreeNode* node, TreeNode* pa = NULL){
        if (!node)
            return;
        parent[node] = pa;
        dfs(node->left, node);
        dfs(node->right, node);
    }
    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        vector<int> ans;
        parent.clear();
        dfs(root);
        queue<pair<TreeNode*, int> > q;
        q.push({ target, 0 });
        set<TreeNode*> visited;
        visited.insert(target);
        while (!q.empty()) {
            pair<TreeNode*, int> temp = q.front();
            q.pop();
            int level = temp.second;
            TreeNode* node = temp.first;
            if (level == k) {
                ans.push_back(node->val);
            }
            if ((node->left && node->left->val != 0) && level + 1 <= k && !visited.count(node->left)) {
                q.push({ node->left, level + 1 });
                visited.insert(node->left);
            }
            if ((node->right && node->right->val != 0) && level + 1 <= k && !visited.count(node->right)){
                q.push({ node->right, level + 1 });
                visited.insert(node->right);
            }
            if (parent[node] != NULL && level + 1 <= k && !visited.count(parent[node])) {
                q.push({ parent[node], level + 1 });
                visited.insert(parent[node]);
            }
        }
        return ans;
    }
};
main(){
    Solution ob;
    vector<int> v = {3,5,1,6,2,0,8,NULL,NULL,7,4};
    TreeNode *root = make_tree(v);
    TreeNode *target = root->left;
    print_vector(ob.distanceK(root, target, 2));
}

入力

{3,5,1,6,2,0,8,NULL,NULL,7,4}

出力

[7, 4, 1]

計算量について

時間計算量・空間計算量はどちらも O(N)(N はノード数)です。DFSによる親の記録でも、BFSによる探索でも、それぞれのノードを高々一度しか処理しないためです。visited 集合の管理によって効率的な探索が実現されています。