Binary Tree Coloring Game LeetCode solution in Python

1
# Definition for a binary tree node.
2
# class TreeNode:
3
#     def __init__(self, val=0, left=None, right=None):
4
#         self.val = val
5
#         self.left = left
6
#         self.right = right
7
class Solution:
8
    def findParent(self, node, par=None):
9
        if node:
10
            self.parent[node.val] = par
11
            self.findParent(node.left, node)
12
            self.findParent(node.right, node)
13

14
    def traverse(self, node, done):
15
        if node:
16
            if node in done:
17
                return 0
18
            done[node] = True
19
            a = self.traverse(self.parent[node.val], done)
20
            b = self.traverse(node.left, done)
21
            c = self.traverse(node.right, done)
22
            return a + b + c + 1
23
        return 0
24

25
    def btreeGameWinningMove(self, root: Optional[TreeNode], n: int, x: int) -> bool:
26
        self.parent = {}
27
        self.findParent(root)
28
        parent = self.parent[x]
29
        node = (
30
            root
31
            if root.val == x
32
            else (
33
                parent.left
34
                if parent and parent.left and parent.left.val == x
35
                else parent.right
36
            )
37
        )
38
        up = self.traverse(parent, {node: True})
39
        left = self.traverse(node.left, {node: True})
40
        right = self.traverse(node.right, {node: True})
41
        return (up > left + right) or (left > up + right) or (right > up + left)