オットセイの経営日誌

データサイエンス系ベンチャーを経営してます。経営のこと、趣味のことつぶやきます。

LeetCode / Invert Binary Tree

初年度のベンチャーはどこもそうなんだと思いますが、特に平日は心身ともに擦り減るので、ベンチャーに休日はないと言えど強制的に休日を作らないとなかなかやっていけません。

と言っても、ただボーッとしてると仕事のことを考えてしまうので、LeetCodeでもやって脳を強制リセットします。

脳の中にコンテナをいくつか作ってうまくオーケストレーションする感じ。あ、また仕事のこと考えてる。。。

https://leetcode.com/problems/invert-binary-tree/

Invert a binary tree.

f:id:mhiro216:20200229203757p:plain

Binary Treeの左右を反転させる問題です。

解法1

Recursiveな解法。

コードにしてみると至って単純で、左の部分木を右に、右の部分木を左に入れ替える処理を再帰的に行います。

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def invertTree(self, root: TreeNode) -> TreeNode:
        if not root:
            return None
        right = self.invertTree(root.right)
        left = self.invertTree(root.left)
        root.left = right
        root.right = left
        return root

時間計算量、空間計算量ともにO(n)です。

解法2

Iterativeな解法。

stackという一時変数を用意して、左右の入れ替えを行う対象を格納していきます。

後は先ほどの解法と同様に、左の部分木を右に、右の部分木を左に入れ替えた上で、入れ替えた部分木をstackに格納し、同様の処理を行っていきます。

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def invertTree(self, root: TreeNode) -> TreeNode:
        stack = [root]
        while stack:
            node = stack.pop()
            if node:
                node.left, node.right = node.right, node.left
                stack += node.left, node.right
        return root

while loop中のpop前のstack, pop後のnodeの中身を見ていくと、以下のようになっています。

len(stack):  1
stack:  [TreeNode{val: 4, left: TreeNode{val: 2, left: TreeNode{val: 1, left: None, right: None}, right: TreeNode{val: 3, left: None, right: None}}, right: TreeNode{val: 7, left: TreeNode{val: 6, left: None, right: None}, right: TreeNode{val: 9, left: None, right: None}}}]
node: TreeNode{val: 4, left: TreeNode{val: 2, left: TreeNode{val: 1, left: None, right: None}, right: TreeNode{val: 3, left: None, right: None}}, right: TreeNode{val: 7, left: TreeNode{val: 6, left: None, right: None}, right: TreeNode{val: 9, left: None, right: None}}}

len(stack):  2
stack:  [TreeNode{val: 7, left: TreeNode{val: 6, left: None, right: None}, right: TreeNode{val: 9, left: None, right: None}}, TreeNode{val: 2, left: TreeNode{val: 1, left: None, right: None}, right: TreeNode{val: 3, left: None, right: None}}]
node: TreeNode{val: 2, left: TreeNode{val: 1, left: None, right: None}, right: TreeNode{val: 3, left: None, right: None}}

len(stack):  3
stack:  [TreeNode{val: 7, left: TreeNode{val: 6, left: None, right: None}, right: TreeNode{val: 9, left: None, right: None}}, TreeNode{val: 3, left: None, right: None}, TreeNode{val: 1, left: None, right: None}]
node: TreeNode{val: 1, left: None, right: None}

len(stack):  4
stack:  [TreeNode{val: 7, left: TreeNode{val: 6, left: None, right: None}, right: TreeNode{val: 9, left: None, right: None}}, TreeNode{val: 3, left: None, right: None}, None, None]
node: None

len(stack):  3
stack:  [TreeNode{val: 7, left: TreeNode{val: 6, left: None, right: None}, right: TreeNode{val: 9, left: None, right: None}}, TreeNode{val: 3, left: None, right: None}, None]
node: None

len(stack):  2
stack:  [TreeNode{val: 7, left: TreeNode{val: 6, left: None, right: None}, right: TreeNode{val: 9, left: None, right: None}}, TreeNode{val: 3, left: None, right: None}]
node: TreeNode{val: 3, left: None, right: None}

len(stack):  3
stack:  [TreeNode{val: 7, left: TreeNode{val: 6, left: None, right: None}, right: TreeNode{val: 9, left: None, right: None}}, None, None]
node: None

len(stack):  2
stack:  [TreeNode{val: 7, left: TreeNode{val: 6, left: None, right: None}, right: TreeNode{val: 9, left: None, right: None}}, None]
node: None

len(stack):  1
stack:  [TreeNode{val: 7, left: TreeNode{val: 6, left: None, right: None}, right: TreeNode{val: 9, left: None, right: None}}]
node: TreeNode{val: 7, left: TreeNode{val: 6, left: None, right: None}, right: TreeNode{val: 9, left: None, right: None}}

len(stack):  2
stack:  [TreeNode{val: 9, left: None, right: None}, TreeNode{val: 6, left: None, right: None}]
node: TreeNode{val: 6, left: None, right: None}

len(stack):  3
stack:  [TreeNode{val: 9, left: None, right: None}, None, None]
node: None

len(stack):  2
stack:  [TreeNode{val: 9, left: None, right: None}, None]
node: None

len(stack):  1
stack:  [TreeNode{val: 9, left: None, right: None}]
node: TreeNode{val: 9, left: None, right: None}

len(stack):  2
stack:  [None, None]
node: None

len(stack):  1
stack:  [None]
node: None

これも時間計算量、空間計算量ともにO(n)です。