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; `
`} ` |

*chevron_right*

*filter_none*

**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.

## Recommended Posts:

- Ways to paint stairs with two colors such that two adjacent are not yellow
- Ways to paint N paintings such that adjacent paintings don't have same colors
- Color tree with minimum colors such that colors of edges incident to a vertex are different
- Number of ways to paint K cells in 3 x N grid such that no P continuous columns are left unpainted
- Count of distinct colors in a subtree of a Colored Tree with given min frequency for Q queries
- Ways to color a skewed tree such that parent and child have different colors
- Check if quantities of 3 distinct colors can be converted to a single color by given merge-pair operations
- Ways to color a 3*N board using 4 colors
- Minimum steps to color the tree with given colors
- Count the nodes of the tree which make a pangram when concatenated with the sub-tree nodes
- Count of all prime weight nodes between given nodes in the given Tree
- Common nodes in the inorder sequence of a tree between given two nodes in O(1) space
- Count root to leaf paths having exactly K distinct nodes in a Binary Tree
- Maximum possible score that can be obtained by constructing a Binary Tree based on given conditions
- Construct XOR tree by Given leaf nodes of Perfect Binary Tree
- Minimum difference between any two weighted nodes in Sum Tree of the given Tree
- Number of ways to form a heap with n distinct integers
- Number of ways to form an array with distinct adjacent elements
- Find number of ways to form sets from N distinct things with no set of size A or B
- Count the number of ways to fill K boxes with N distinct items

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.