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

• Difficulty Level : Hard
• Last Updated : 26 Sep, 2021

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 :
We have three color, say red, blue and green. we can paint in the following ways.

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 :

CPP

 // C++ Implementation of above approach#include using namespace std;const int maxx = 1e5;vector 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 permutations ( 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);        }        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.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 Codeint 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