Program to check whether we can color a tree where no adjacent nodes have the same color or not in python

Suppose we have a binary tree where the value of each node represents its color. There are at most 2 colors in a tree. We have to check whether it is possible to swap the colors of the nodes any number of times so that no two connected nodes have the same color.

So, if the input is like:

then the output will be True as we can rearrange colors to avoid adjacent nodes having the same color.

Algorithm

To solve this, we will follow these steps ?

• Create two maps: colors to count nodes with each color value, and prop to count nodes at each level parity
• Define a DFS function that traverses the tree while tracking level parity (odd/even levels)
• For each node, increment counters for its color and current level parity
• Recursively traverse left and right children with flipped parity
• Check if color distribution matches level parity distribution

Implementation

from collections import defaultdict

class TreeNode:
    def __init__(self, data, left=None, right=None):
        self.val = data
        self.left = left
        self.right = right

class Solution:
    def solve(self, root):
        colors = defaultdict(int)
        prop = defaultdict(int)

        def dfs(node, flag=True):
            if not node:
                return
            colors[node.val] += 1
            prop[flag] += 1
            dfs(node.left, not flag)
            dfs(node.right, not flag)

        dfs(root)
        return set(colors.values()) == set(prop.values())

# Create the binary tree
ob = Solution()
root = TreeNode(2)
root.left = TreeNode(2)
root.right = TreeNode(2)
root.right.left = TreeNode(1)
root.right.right = TreeNode(1)
root.right.left.right = TreeNode(1)

print(ob.solve(root))

True

How It Works

The algorithm works by recognizing that in a properly colored tree:

• Nodes at even levels should have one color
• Nodes at odd levels should have the other color
• The count of each color must match the count of nodes at corresponding level parities

In our example, we have 4 nodes with color value 2 and 2 nodes with color value 1. After DFS traversal with alternating flags, if the distribution of colors matches the distribution of level parities, then we can successfully color the tree.

Example with Different Input

# Test with a tree that cannot be properly colored
ob = Solution()
root2 = TreeNode(1)
root2.left = TreeNode(1)
root2.right = TreeNode(1)
root2.left.left = TreeNode(1)

print(ob.solve(root2))

False

Conclusion

This algorithm efficiently checks if a binary tree can be properly 2-colored by comparing color distribution with level parity distribution. The time complexity is O(n) where n is the number of nodes, making it an optimal solution for this tree coloring problem.