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 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 <bits/stdc++.h> ` `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[1]); ` ` ` `} ` ` ` `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[2].push_back(1); ` ` ` `tree[1].push_back(2); ` ` ` `tree[3].push_back(2); ` ` ` `tree[2].push_back(3); ` ` ` `degree_of_node[2]++; ` ` ` `degree_of_node[1]++; ` ` ` `degree_of_node[3]++; ` ` ` `degree_of_node[2]++; ` `} ` ` ` `// 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.