Problem statement

https://binarysearch.com/problems/Tree-Coloring/

Solution

Any tree is bipartite graph. So, we need to check that we have enough colors for each part.

Complexity

It is O(n) for time and O(h) for space.

Code

class Solution:
    def solve(self, root):
        cols = Counter()
        deps = [0, 0]
        def dfs(node, d):
            if not node: return
            cols[node.val] += 1
            deps[d%2] += 1
            dfs(node.left, d + 1)
            dfs(node.right, d + 1)
        
        dfs(root, 0)
        A, B = sorted(list(cols.values())), sorted(deps)
        A = [0]*(2 - len(A)) + A
        return A == B