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.