# Number of ways to paint a tree of N nodes with K distinct colors with given conditions

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 1 Node 2 Node 3
Red Blue Green
Red Green Blue
Blue Red Green
Blue Green Red
Green Red Blue
Green Blue Red

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)

My Personal Notes arrow_drop_up Check out this Author's contributed articles.

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.

Article Tags :
Practice Tags :

2

Please write to us at contribute@geeksforgeeks.org to report any issue with the above content.