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)
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 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<Integer>[] 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 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 |
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 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(); |
6
Time Complexity : O(N)