C++で指定された値を持つ葉ノードを削除するアルゴリズム

問題の概要

二分木と整数 target が与えられたとき、値が target と一致するすべての葉ノードを削除することを考えます。ここで重要なのは、葉ノードを削除した結果、その親ノードが新たに葉ノードになり、かつその値が target と一致する場合には、その親ノードも同様に削除しなければならないという点です。この操作は、削除できるノードがなくなるまで繰り返し行います。

例えば、下図のような二分木があり、target が 2 の場合、最終的な木は次のようになります。

解法のアプローチ

この問題は、再帰を用いた後順（ボトムアップ）処理によって効率的に解くことができます。具体的な手順は以下の通りです。

• ルートノードと target を引数に取る再帰メソッド remLeaf() を定義します。

• ルートが null の場合は null を返します。

• left := remLeaf(ルートの左の子, target) を実行します。

• right := remLeaf(ルートの右の子, target) を実行します。

• left と right がどちらも null であり、かつルートの値が target と一致する場合は null を返します（このノードを削除）。

• ルートの左の子 := left、右の子 := right を設定します。

• ルートを返します。

子ノードから先に処理を行うことで、葉が削除されて新たに葉になった親ノードも自動的に判定・削除される仕組みになっています。

C++による実装例

以下の実装を見ると、より理解が深まるでしょう。

#include <bits/stdc++.h>
using namespace std;
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;
}
void tree_level_trav(TreeNode*root){
    if (root == NULL) return;
    cout << "[";
    queue<TreeNode *> q;
    TreeNode *curr;
    q.push(root);
    q.push(NULL);
    while (q.size() > 1) {
        curr = q.front();
        q.pop();
        if (curr == NULL){
            q.push(NULL);
        } else {
            if(curr->left)
                q.push(curr->left);
            if(curr->right)
                q.push(curr->right);
            if(curr->val == 0 || curr == NULL){
                cout << "null" << ", ";
            } else {
                cout << curr->val << ", ";
            }
        }
    }
    cout << "]"<<endl;
}
class Solution {
    public:
        TreeNode* removeLeafNodes(TreeNode* root, int target) {
            if(!root || root->val == 0) return NULL;
            TreeNode* left = removeLeafNodes(root->left, target);
            TreeNode* right = removeLeafNodes(root->right, target);
            if(!left && !right && root->val == target){
                return NULL;
            }
            root->left = left;
            root->right = right;
            return root;
        }
};
main() {
    vector<int> v1 = {1,2,3,2,NULL,2,4};
    TreeNode *root = make_tree(v1);
    Solution ob;
    tree_level_trav(ob.removeLeafNodes(root, 2));
}

入力

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

出力

[1, 3, 4]

計算量について

このアルゴリズムは各ノードを一度だけ訪問するため、時間計算量は O(N)（N はノード数）、空間計算量は再帰呼び出しの深さに依存し、最悪の場合 O(N) となります。木が平衡な場合は O(log N) に抑えられます。