Skip to content
Related Articles

Related Articles

Number of ways to paint a tree of N nodes with K distinct colors with given conditions
  • Difficulty Level : Medium
  • Last Updated : 28 Jun, 2019

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 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 :




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

My Personal Notes arrow_drop_up
Recommended Articles
Page :