C++で葉ノードから距離kにあるすべてのノードを出力する方法

問題概要

この問題では、二分木と数値Kが与えられ、葉ノードから距離Kにあるすべてのノードを出力することが求められます。

二分木（Binary Tree）とは、各ノードが最大2つの子ノード（1つ・2つ・または0個）を持つ特別な木構造のことです。

葉ノード（Leaf Node）とは、二分木の末端に位置するノードを指します。

この問題における「葉ノードからの距離」とは、葉ノードよりも上位のレベルに位置するノードを意味します。たとえば、レベル4にある葉ノードから距離2のノードは、レベル2に存在することになります。

具体例で理解しよう

次の図のような二分木を例に考えてみましょう。

K = 2 の場合、出力：6 9

解法のアプローチ

この問題を解くには、まず木全体を走査します。葉ノードに到達するまで、すべての親ノード（祖先ノードとも呼ばれます）の情報をレベルごとに保存していきます。そして、葉ノードから距離kだけ離れた祖先ノードを出力します。

走査中に同じノードが複数回出力されるのを防ぐため、訪問済みノードのマーク付けが重要です。ここではブール型（bool）の配列を使用して、すでに出力したノードを管理します。

このアルゴリズムは木の走査のみを行うため、計算量はノード数nに比例します。

時間計算量：O(n)

実装例

上記のロジックを実装したC++プログラムは以下の通りです。

#include <iostream>
using namespace std;
#define MAX_HEIGHT 10000

struct Node {
    int key;
    Node *left, *right;
};

Node* insertNode(int key){
    Node* node = new Node;
    node->key = key;
    node->left = node->right = NULL;
    return (node);
}

void nodesKatDistance(Node* node, int path[], bool visited[], int pathLen, int k){
    if (node == NULL) return;
    path[pathLen] = node->key;
    visited[pathLen] = false;
    pathLen++;
    if (node->left == NULL && node->right == NULL
        && pathLen-k-1 >= 0
        && visited[pathLen-k-1] == false){
        cout << path[pathLen-k-1] << "\t";
        visited[pathLen-k-1] = true;
        return;
    }
    nodesKatDistance(node->left, path, visited, pathLen, k);
    nodesKatDistance(node->right, path, visited, pathLen, k);
}

void printNodes(Node* node, int k){
    int path[MAX_HEIGHT];
    bool visited[MAX_HEIGHT] = {false};
    nodesKatDistance(node, path, visited, 0, k);
}

int main(){
    Node *root = insertNode(6);
    root->left = insertNode(3);
    root->right = insertNode(9);
    root->left->right = insertNode(4);
    root->right->left = insertNode(8);
    root->right->right = insertNode(10);
    root->right->left->left = insertNode(5);
    root->right->left->right = insertNode(1);
    int k = 2;
    cout << "All nodes at distance " << k << " from leaf node are:\n";
    printNodes(root, k);
    return 0;
}

出力結果

葉ノードから距離2にあるすべてのノードの出力は以下の通りです。

6 9

まとめ

本記事では、二分木において葉ノードから距離kにあるノードをすべて出力する方法を解説しました。再帰的な深さ優先探索で各経路を記録し、ブール配列による訪問済み管理で重複を排除するのがポイントです。計算量はO(n)と効率的で、木構造の基本的な走査テクニックを学ぶのに最適な問題といえます。