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 1void 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 conditionsint 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 treevoid 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 Codeint main(){ int N = 3, K = 3; MakeTree(); int Count = NoOfWays(N, K); cout << Count << "\n"; return 0;} |
6
Time Complexity : O(N)
