Skip to content
Related Articles

Related Articles

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

View Discussion
Improve Article
Save Article
  • Difficulty Level : Hard
  • Last Updated : 26 Sep, 2021
View Discussion
Improve Article
Save Article

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. 

 

Node 1Node 2Node 3
RedBlueGreen
RedGreenBlue
BlueRedGreen
BlueGreenRed
GreenRedBlue
GreenBlueRed

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 : 
 

CPP




// 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 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 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)
 


My Personal Notes arrow_drop_up
Recommended Articles
Page :

Start Your Coding Journey Now!