C++で二分木をリンクリストにフラット化（平坦化）する方法

二分木が与えられたとき、それをその場（in-place）でリンクリストへフラット化（平坦化）することを考えます。具体的には、すべてのノードを右ポインタで連結し、左ポインタを null にした、一本の連結リストのような構造へ変換します。

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

これをフラット化すると、出力は次のようになります。

アルゴリズムの手順

この問題は、逆後順走査（右 → 左 → 根）を利用することで効率的に解けます。手順は以下の通りです。

• prev を null で初期化します。
• ルートを引数にとる再帰関数 solve() を定義します。
• root が null の場合は、そのまま戻ります。
• まず root の右部分木に対して solve() を呼び出します。
• 次に root の左部分木に対して solve() を呼び出します。
• root->right を prev に、root->left を null に設定します。
• prev を root に更新します。

この走査順序により、処理済みのノードが常に「すでにフラット化されたリストの先頭」として prev に保持され、各ノードの right ポインタを順につなぎ直していくことができます。

C++での実装例

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

#include <bits/stdc++.h>
using namespace std;
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:
    TreeNode* prev = NULL;
    void flatten(TreeNode* root) {
        if(!root) return;
        flatten(root->right);
        flatten(root->left);
        root->right = prev;
        root->left = NULL;
        prev = root;
    }
};
main(){
    vector<int> v = {1,2,5,3,4};
    TreeNode *root = make_tree(v);
    Solution ob;
    (ob.flatten(root));
    TreeNode *ptr = root;
    while(ptr != NULL && ptr->val != 0){
        cout << ptr->val << ", ";
        ptr = ptr->right;
    }
}

入力

[1,2,5,3,4]

出力

1, 2, 3, 4, 5,

計算量

時間計算量： O(n) ― 各ノードを一度ずつ訪問するだけです。
空間計算量： O(h) ― 再帰スタックの深さ分が必要です（h は木の高さ）。平衡な木では O(log n)、最悪ケース（線形リスト状の木）では O(n) となります。