Pythonで二分木の葉ノードを新しいルートに変更するプログラムの実装方法

二分木と、その葉(リーフ)に位置する1つのノードが与えられたとしましょう。ここでの課題は、その葉ノードを二分木の新しいルート(根)ノードへと変更することです。この操作は、次の2つのルールに従って行います。

• 左の子の移動: ノードに左の子が存在する場合、その子は右側へ移動します。
• 親の移動: ノードの親は、そのノードの左の子になります。この処理の過程で、元の親ノードからそのノードへのリンクは切断(null)されるため、親ノードは子を1つだけ持つ状態になります。

今回扱うツリーのノード構造は以下の通りです。

TreeNode:
    data: <整数>
    left: <TreeNodeへのポインタ>
    right: <TreeNodeへのポインタ>
    parent: <TreeNodeへのポインタ>

処理の結果として、変換後のツリーのルートを返す必要があります。

具体例

例えば、次のような二分木が与えられたとします。

ここで新しいルートを「8」に変更すると、変換後のツリーの中間順(inorder)走査の結果は次のようになります。

2, 3, 4, 5, 7, 6, 8,

ご覧の通り、新しいツリーのルートノードは8となっています。

解決アプローチ

この問題を解くために、再帰的なヘルパー関数 helper() を定義し、以下の手順で処理を進めます。

1. helper(node, new_par) 関数を定義する。
   • node が root と一致する場合:
     • node の親を new_par に設定する。
     • node の左の子が new_par と一致する場合は、左の子を null にする。
     • node の右の子が new_par と一致する場合は、右の子を null にする。
     • root を返す。
   • node の左の子が null でない場合、node の右の子を現在の左の子に置き換える。
   • node の親の左の子が node と一致する場合、親の左の子を null にする。
   • node の左の子を helper(node の親, node) の戻り値に設定する。
   • node の親を new_par に更新する。
   • node を返す。
2. helper(leaf, None) を呼び出した結果を返す。

それでは、実際の実装例を見ながら理解を深めていきましょう。

実装例(Python)

import collections
class TreeNode:
    def __init__(self, data, left = None, right = None, parent = None):
        self.data = data
        self.left = left
        self.right = right
        self.parent = parent
    def insert(temp,data):
        que = []
        que.append(temp)
        while (len(que)):
            temp = que[0]
            que.pop(0)
            if (not temp.left):
                if data is not None:
                    temp.left = TreeNode(data, parent = temp)
                else:
                    temp.left = TreeNode(0, parent = temp)
                break
            else:
                que.append(temp.left)
            if (not temp.right):
                if data is not None:
                    temp.right = TreeNode(data, parent = temp)
                else:
                    temp.right = TreeNode(0, parent = temp)
                break
            else:
                que.append(temp.right)
def make_tree(elements):
    Tree = TreeNode(elements[0])
    for element in elements[1:]:
        insert(Tree, element)
    return Tree
def search_node(root, element):
    if (root == None):
        return None
    if (root.data == element):
        return root
    res1 = search_node(root.left, element)
    if res1:
        return res1
    res2 = search_node(root.right, element)
    return res2
def print_tree(root):
    if root is not None:
        print_tree(root.left)
        print(root.data, end = ', ')
        print_tree(root.right)
def solve(root, leaf):
    def helper(node, new_par):
        if node == root:
            node.parent = new_par
            if node.left == new_par:
                node.left = None
            if node.right == new_par:
                node.right = None
            return root
        if node.left:
            node.right = node.left
        if node.parent.left == node:
            node.parent.left = None
        node.left = helper(node.parent, node)
        node.parent = new_par
        return node
    return helper(leaf, None)
root = make_tree([5, 3, 7, 2, 4, 6, 8])
root = solve(root, search_node(root, 8))
print_tree(root)

入力

root = make_tree([5, 3, 7, 2, 4, 6, 8])
root = solve(root, search_node(root, 8))

出力

2, 3, 4, 5, 7, 6, 8,