Related Articles
Number of ways to paint a tree of N nodes with K distinct colors with given conditions
• Difficulty Level : Medium
• Last Updated : 28 Jun, 2019

Given a tree with N nodes and a number K. Paint every node of the tree in one of the K available colors.
Count and return the number of ways of painting the tree such that any two nodes that are at distance 1 or 2 are painted in different colors.

Examples:The first line of the input contains two integer N and K.
The next line contains an array of pairs. Each pair (x, y) denotes an undirected edge between x and y.

Input : N = 3 K = 3
Tree = { (2, 1), (3, 2) }
Output : 6
We have three color, say red, blue and green. we can paint in the following ways.

Node 1Node 2Node 3
RedBlueGreen
RedGreenBlue
BlueRedGreen
BlueGreenRed
GreenRedBlue
GreenBlueRed

Thus 6 is the answer.
Input : N = 5 K = 6
Tree= { (1, 2), (5, 1), (3, 1), (4, 2) }
Output :48

Approach :
Let’s root the tree at node 1, and then we paint it starting with the root moving down to the leaves. For the root we can paint it with k available colors. If the root has x children we can paint it with k-1 P x ways, that is
(k-1)!/(k-1-x)!. Because each child has to use a distinct color, and they all should be different from the color used for the root.

Now for the remaining nodes, we paint all the sons of a particular node v at once. Their colors have to be distinct and different from the color used for v and v’s father. So if v has x sons, we can paint them in k-2 P x ways

Below is the implementation of above approach :

 `// C++ Implementation of above approach``#include ``using` `namespace` `std;``const` `int` `maxx = 1e5;``vector<``int``> tree[maxx];``int` `degree_of_node[maxx], parent_of_node[maxx],``    ``child_of_node[maxx], flag = -1;`` ` `// Function to calculate number of children``// of every node in a tree with root 1``void` `dfs(``int` `current, ``int` `parent)``{``    ``parent_of_node[current] = parent;``    ``for` `(``int``& child : tree[current]) {`` ` `        ``// If current and parent are same we have``        ``// already visited it, so no need to visit again``        ``if` `(child == parent)``            ``return``;``        ``dfs(child, current);``    ``}`` ` `    ``// If the current node is a leaf node``    ``if` `(degree_of_node[current] == 1 && current != 1) {`` ` `        ``// For leaf nodes there will be no child.``        ``child_of_node[current] = 0;``        ``return``;``    ``}`` ` `    ``// Gives the total child of current node``    ``int` `total_child = 0;``    ``for` `(``auto``& child : tree[current]) {``        ``if` `(child == parent)``            ``return``;``        ``else``            ``++total_child;``    ``}``    ``child_of_node[current] = total_child;``    ``return``;``}`` ` `// Function to calculate permuations ( nPr )``int` `find_nPr(``int` `N, ``int` `R)``{``    ``if` `(R > N) {``        ``flag = 0;``        ``return` `0;``    ``}``    ``int` `total = 1;``    ``for` `(``int` `i = N - R + 1; i <= N; ++i) {``        ``total = total * i;``    ``}``    ``return` `total;``}`` ` `// Function to calculate the number of ways``// to paint the tree according to given conditions``int` `NoOfWays(``int` `Nodes, ``int` `colors)``{`` ` `    ``// Do dfs to find parent and child of a node,``    ``// we root the tree at node 1.``    ``dfs(1, -1);`` ` `    ``// Now start iterating for all nodes of``    ``// the tree and count the number of ways to``    ``// paint its children and node itself``    ``int` `ways = 0;``    ``for` `(``int` `i = 1; i <= Nodes; ++i) {`` ` `        ``// If the current node is root node, then``        ``// we have total of K ways to paint it and``        ``// (k-1)P(x) to paint its child``        ``if` `(i == 1) {``            ``ways = ways + colors * ``                   ``find_nPr(colors - 1, child_of_node);``        ``}``        ``else` `{`` ` `            ``// For other remaining nodes which are not``            ``// leaf nodes we have (k-2)P(x) to paint``            ``// its children, we will not take into``            ``// consideration of current node``            ``// since we already painted it.``            ``if` `(degree_of_node[i] == 1) {``                ``continue``;``            ``}``            ``else` `{``                ``ways = ways * ``                ``find_nPr(colors - 2, child_of_node[i]);``            ``}``        ``}``    ``}``    ``return` `ways;``}`` ` `// Function to build the tree``void` `MakeTree()``{`` ` `    ``tree.push_back(1);``    ``tree.push_back(2);``    ``tree.push_back(2);``    ``tree.push_back(3);``    ``degree_of_node++;``    ``degree_of_node++;``    ``degree_of_node++;``    ``degree_of_node++;``}`` ` `// Driver Code``int` `main()``{``    ``int` `N = 3, K = 3;``    ``MakeTree();``    ``int` `Count = NoOfWays(N, K);``    ``cout << Count << ``"\n"``;``    ``return` `0;``}`
Output:
```6
```

Time Complexity : O(N)

Attention reader! Don’t stop learning now. Get hold of all the important DSA concepts with the DSA Self Paced Course at a student-friendly price and become industry ready.

My Personal Notes arrow_drop_up