leetcode-practice-python/0226__invert_binary_tree.py at master · Philosoft/leetcode-practice-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
from typing import Optional
from unittest import TestCase

from lib.TreeNode import TreeNode, build_tree


class Solution(TestCase):
    def test_example_3(self):
        root = self.invertTree(build_tree([4, 2, 7, 1, 3, 6, 9]))
        left, right = root.left, root.right

        self.assertEqual(7, left.val)
        self.assertEqual(2, right.val)

        self.assertEqual(9, left.left.val)
        self.assertEqual(1, right.right.val)

    def test_example_2(self):
        root = build_tree([2, 1, 3])
        root = self.invertTree(root)

        self.assertEqual(3, root.left.val)
        self.assertEqual(1, root.right.val)

    def test_example_3(self):
        self.assertIsNone(self.invertTree(None))

    def invertTree(self, root: Optional[TreeNode]) -> Optional[TreeNode]:
        if not root:
            return None

        root.left, root.right = self.invertTree(root.right), self.invertTree(root.left)

        return root