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

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)
        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;
    // Gives the total child of current node
    int total_child = 0;
    for (auto& child : tree[current]) {
        if (child == parent)
    child_of_node[current] = total_child;
// 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) {
            else {
                ways = ways * 
                find_nPr(colors - 2, child_of_node[i]);
    return ways;
// Function to build the tree
void MakeTree()
// Driver Code
int main()
    int N = 3, K = 3;
    int Count = NoOfWays(N, K);
    cout << Count << "\n";
    return 0;


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:

Check out this Author's contributed articles.

If you like GeeksforGeeks and would like to contribute, you can also write an article using or mail your article to 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.

Article Tags :
Practice Tags :