Tree Coloring - Amazon Top Interview Questions

                        Solution in C++ :

/**
 * class Tree {
 *     public:
 *         int val;
 *         Tree *left;
 *         Tree *right;
 * };
 */

int dfs(Tree *root, int color) {
    if (!root) return 0;
    return color + dfs(root->left, 1 - color) + dfs(root->right, 1 - color);
}

void countNodes(Tree *root, int firstColor, int &firstCount, int &secondCount) {
    if (!root) return;

    if (root->val == firstColor)
        ++firstCount;
    else
        ++secondCount;

    countNodes(root->left, firstColor, firstCount, secondCount);
    countNodes(root->right, firstColor, firstCount, secondCount);
}

bool solve(Tree *root) {
    if (!root) return true;

    int req = dfs(root, 1), firstColor = root->val, firstCount = 0, secondCount = 0;
    countNodes(root, firstColor, firstCount, secondCount);
    return req == firstCount || req == secondCount;
}
                    

                        Solution in Python : 
                            
# class Tree:
#     def __init__(self, val, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def solve(self, root):
        if not root:
            return True
        evens = 0  # how many nodes are in even layers
        tot = 0  # how many total nodes there are
        firstCnt = 0  # how many nodes have the "first" color
        first = root.val  # we assign the root node's color to be the first color
        q = collections.deque([(root, 0)])
        while q:
            curr, lvl = q.popleft()
            tot += 1
            if curr.val == first:
                firstCnt += 1
            if lvl % 2:
                evens += 1
            if curr.left:
                q.append((curr.left, lvl + 1))
            if curr.right:
                q.append((curr.right, lvl + 1))
        return evens in [
            firstCnt,
            tot - firstCnt,
        ]  # the even layers can get the first or second color