Open In App

# 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 :
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<``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;``}`

## Java

 `//Java code for the above approach``import` `java.util.ArrayList;``import` `java.util.List;` `class` `GFG {``    ``static` `int``[] degree_of_node = ``new` `int``[(``int``)1e5];``    ``static` `int``[] parent_of_node = ``new` `int``[(``int``)1e5];``    ``static` `int``[] child_of_node = ``new` `int``[(``int``)1e5];``    ``static` `List[] tree = ``new` `ArrayList[(``int``)1e5];``    ``static` `int` `flag = -``1``;` `    ``// Function to calculate number of children``    ``// of every node in a tree with root 1``    ``static` `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` `(``int` `child : tree[current]) {``            ``if` `(child == parent)``                ``return``;``            ``else``                ``++total_child;``        ``}``        ``child_of_node[current] = total_child;``        ``return``;``    ``}` `    ``// Function to calculate permutations ( nPr )``    ``static` `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``    ``static` `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``    ``static` `void` `MakeTree()``    ``{``        ``for` `(``int` `i = ``0``; i < (``int``)1e5; i++) {``            ``tree[i] = ``new` `ArrayList<>();``        ``}` `        ``tree[``2``].add(``1``);``        ``tree[``1``].add(``2``);``        ``tree[``3``].add(``2``);``        ``tree[``2``].add(``3``);``        ``degree_of_node[``2``]++;``        ``degree_of_node[``1``]++;``        ``degree_of_node[``3``]++;``        ``degree_of_node[``2``]++;``    ``}` `    ``// Driver Code``    ``public` `static` `void` `main(String[] args)``    ``{``        ``int` `N = ``3``, K = ``3``;``        ``MakeTree();``        ``int` `Count = NoOfWays(N, K);``        ``System.out.println(Count);``    ``}``}`

## Python3

 `#Python3 code for the above approach``from` `typing ``import` `List` `def` `dfs(tree: ``List``[``List``[``int``]], degree_of_node: ``List``[``int``], parent_of_node: ``List``[``int``], child_of_node: ``List``[``int``], current: ``int``, parent: ``int``) ``-``> ``None``:``    ``parent_of_node[current] ``=` `parent``    ``for` `child ``in` `tree[current]:``        ``# If current and parent are same we have``        ``# already visited it, so no need to visit again``        ``if` `child ``=``=` `parent:``            ``return``        ``dfs(tree, degree_of_node, parent_of_node, child_of_node, child, current)` `    ``# If the current node is a leaf node``    ``if` `degree_of_node[current] ``=``=` `1` `and` `current !``=` `1``:``        ``# For leaf nodes there will be no child.``        ``child_of_node[current] ``=` `0``        ``return``    ``# Gives the total child of current node``    ``total_child ``=` `0``    ``for` `child ``in` `tree[current]:``        ``if` `child ``=``=` `parent:``            ``return``        ``else``:``            ``total_child ``+``=` `1``    ``child_of_node[current] ``=` `total_child` `# Function to calculate permutations ( nPr )``def` `find_nPr(N: ``int``, R: ``int``) ``-``> ``int``:``    ``if` `R > N:``        ``return` `0``    ``total ``=` `1``    ``for` `i ``in` `range``(N ``-` `R ``+` `1``, N ``+` `1``):``        ``total ``*``=` `i``    ``return` `total` `# Function to calculate the number of ways``# to paint the tree according to given conditions``def` `NoOfWays(Nodes: ``int``, colors: ``int``) ``-``> ``int``:``    ``# Do dfs to find parent and child of a node,``    ``# we root the tree at node 1.``    ``dfs(tree, degree_of_node, parent_of_node, child_of_node, ``1``, ``-``1``)` `    ``# Now start iterating for all nodes of``    ``# the tree and count the number of ways to``    ``# paint its children and node itself``    ``ways ``=` `0``    ``for` `i ``in` `range``(``1``, Nodes ``+` `1``):``        ``# 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 ``+``=` `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 ``*``=` `find_nPr(colors ``-` `2``, child_of_node[i])``    ``return` `ways` `# Function to build the tree``def` `MakeTree() ``-``> ``None``:``    ``tree[``2``].append(``1``)``    ``tree[``1``].append(``2``)``    ``tree[``3``].append(``2``)``    ``tree[``2``].append(``3``)``    ``degree_of_node[``2``] ``+``=` `1``    ``degree_of_node[``1``] ``+``=` `1``    ``degree_of_node[``3``] ``+``=` `1``    ``degree_of_node[``2``] ``+``=` `1``    ` `# Driver Code``if` `__name__ ``=``=` `"__main__"``:``    ``N ``=` `3``    ``K ``=` `3``    ``tree ``=` `[[] ``for` `_ ``in` `range``(N ``+` `1``)]``    ``degree_of_node ``=` `[``0` `for` `_ ``in` `range``(N ``+` `1``)]``    ``parent_of_node ``=` `[``0` `for` `_ ``in` `range``(N ``+` `1``)]``    ``child_of_node ``=` `[``0` `for` `_ ``in` `range``(N ``+` `1``)]``    ``MakeTree()``    ``Count ``=` `NoOfWays(N, K)``    ``print``(Count)``    ` `#This code is contributed by Potta Lokesh`

## C#

 `using` `System;``using` `System.Collections.Generic;` `class` `GFG``{``  ` `  ``// Variables to store the degree of each node,``  ``// the parent of each node, and the number of children``  ``// of each node in the tree rooted at node 1``  ``static` `int``[] degree_of_node = ``new` `int``[100000];``  ``static` `int``[] parent_of_node = ``new` `int``[100000];``  ``static` `int``[] child_of_node = ``new` `int``[100000];``  ``static` `List<``int``>[] tree = ``new` `List<``int``>[100000];``  ``static` `int` `flag = -1;` `  ``// Variables to store the degree of each node,``  ``// the parent of each node, and the number of children``  ``// of each node in the tree rooted at node 1``  ``static` `void` `dfs(``int` `current, ``int` `parent)``  ``{``    ``parent_of_node[current] = parent;``    ``foreach` `(``int` `child ``in` `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)``    ``{``      ``child_of_node[current] = 0;``      ``return``;``    ``}` `    ``// Gives the total child of current node``    ``int` `total_child = 0;``    ``foreach` `(``int` `child ``in` `tree[current])``    ``{``      ``if` `(child == parent)``        ``return``;``      ``else``        ``++total_child;``    ``}``    ``child_of_node[current] = total_child;``    ``return``;``  ``}` `  ``// Function to calculate permutations ( nPr )``  ``static` `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``  ``static` `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]);``      ``}` `      ``// 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.``      ``else``      ``{``        ``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``  ``static` `void` `MakeTree()``  ``{``    ``for` `(``int` `i = 0; i < 100000; i++)``    ``{``      ``tree[i] = ``new` `List<``int``>();``    ``}` `    ``tree[2].Add(1);``    ``tree[1].Add(2);``    ``tree[3].Add(2);``    ``tree[2].Add(3);``    ``degree_of_node[2]++;``    ``degree_of_node[1]++;``    ``degree_of_node[3]++;``    ``degree_of_node[2]++;``  ``}``  ` `  ``//Driver code``  ``static` `void` `Main(``string``[] args)``  ``{``    ``int` `N = 3, K = 3;``    ``MakeTree();``    ``int` `Count = NoOfWays(N, K);``    ``Console.WriteLine(Count);``  ``}``}`

## Javascript

 `// JavaScript code for the above approach``let degree_of_node = ``new` `Array(1e5).fill(0);``let parent_of_node = ``new` `Array(1e5).fill(0);``let child_of_node = ``new` `Array(1e5).fill(0);``let tree = ``new` `Array(1e5);``for` `(let i = 0; i < 1e5; i++) {``    ``tree[i] = [];``}``let flag = -1;` `// Function to calculate number of children``// of every node in a tree with root 1``function` `dfs(current, parent) {``    ``parent_of_node[current] = parent;``    ``for` `(let child of 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``  ``let total_child = 0;``  ``for` `(let child of tree[current]) {``      ``if` `(child === parent)``          ``return``;``      ``else``          ``++total_child;``  ``}``  ``child_of_node[current] = total_child;``  ``return``;` `}` `// Function to calculate permutations ( nPr )``function` `find_nPr(N, R) {``    ``if` `(R > N) {``    ``flag = 0;``    ``return` `0;``    ``}``    ``let total = 1;``    ``for` `(let 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``function` `NoOfWays(Nodes, 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``    ``let ways = 0;``    ``for` `(let 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``function` `MakeTree() {``  ``tree[2].push(1);``  ``tree[1].push(2);``  ``tree[3].push(2);``  ``tree[2].push(3);``  ``degree_of_node[2]++;``  ``degree_of_node[1]++;``  ``degree_of_node[3]++;``  ``degree_of_node[2]++;``}``function` `main() {``  ``const N = 3, K = 3;``  ``MakeTree();``  ``const Count = NoOfWays(N, K);``  ``console.log(Count);``}` `main();`

Output:

`6`

Time Complexity : O(N)