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 : 

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;
}

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

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