Print node whose each neighboring Tree has all nodes of same color

Given a tree with N nodes numbered from 1 to N and N – 1 edge and array colours[] where colours[i] denote the color of i^th Node. The task is to find a node such that each neighboring tree connected to this node consists of the same colored nodes. If no such node exists then print -1.

Input: N = 8, colours[] = {1, 1, 1, 1, 1, 2, 1, 3} edges = {(1, 2) (1, 3) (2, 4) (2, 7) (3, 5) (3, 6) (6, 8)}

Visualizing the tree

Output: 6
Explanation:
Consider the node 6, it has 2 trees connected to it. One of them is rooted at 3 and the other is rooted at 8. Clearly, the tree rooted at 3 has nodes of same color and the tree rooted at 8 has only one node. Therefore, node 6 is one such node.

Input: N = 4, colors[] = {1, 2, 3, 4}, edges = {(1, 3) (1, 2 ) (2, 4)}
Output: -1
Explanation:
There is no such node.

Approach: The idea is to check if all the nodes have the same color, then any node can be the root. Otherwise, pick any two nodes which are adjacent to each other and have different colors and check the subtrees of these nodes by performing DFS. If any of these nodes satisfy the condition then that node can be the root. If none of these two nodes satisfies the condition then no such root exists and print -1.

1. Traverse over the tree and find the first two different color nodes which are adjacent to each other let’s say root1 and root2. If no such nodes are found, then all nodes are of the same color and any node can then be taken as root.
2. Check if of all nodes of each subtree have the same color by considering root1 as the root of the tree. If the condition is satisfied then root1 is the answer.
3. Repeat step 2 for root2 if root1 does not satisfy the condition.
4. If root2 does not satisfy the condition then no such root exists and the output is -1.

Below is the implementation of the above approach:

C++

// C++ program for the above approach
#include <bits/stdc++.h>
using namespace std;

const int NN = 1e5 + 5;
// Vector to store the tree
vector<int> G[NN];

// Function to perform dfs
void dfs(int node, int parent,
         bool& check,
         int current_colour,
         int* colours)
{
    // Check is assigned to false if either it
    // is already false or the current_colour
    // is not same as the node colour
    check = check
            && (colours[node] == current_colour);

    // Iterate over the neighbours of node
    for (auto a : G[node]) {

        // If the neighbour is
        // not the parent node
        if (a != parent) {

            // call the function
            // for the neighbour
            dfs(a, node, check,
                current_colour,
                colours);
        }
    }
}
// Function to check whether all the
// nodes in each subtree of the given
// node have same colour
bool checkPossibility(
    int root, int* colours)
{

    // Initialise the boolean answer
    bool ans = true;
    // Iterate over the neighbours
    // of selected root
    for (auto a : G[root]) {

        // Initialise the colour
        // for this subtree
        // as the colour of
        // first neighbour
        int current_colour = colours[a];

        // Variable to check
        // condition of same
        // colour for each subtree
        bool check = true;

        // dfs function call
        dfs(a, root, check,
            current_colour, colours);

        // Check if any one subtree
        // does not have all
        // nodes of same colour
        // then ans will become false

        ans = ans && check;
    }

    // Return the answer
    return ans;
}

// Function to add edges to the tree
void addedge(int x, int y)
{
    // y is added as a neighbour of x
    G[x].push_back(y);

    // x is added as a neighbour of y
    G[y].push_back(x);
}

// Function to find the node
void solve(int* colours, int N)
{
    // Initialise root1 as -1
    int root1 = -1;

    // Initialise root2 as -1
    int root2 = -1;

    // Find the first two nodes of
    // different colour which are adjacent
    // to each other
    for (int i = 1; i <= N; i++) {
        for (auto a : G[i]) {
            if (colours[a] != colours[i]) {
                root1 = a;
                root2 = i;
                break;
            }
        }
    }

    // If no two nodes of different
    // colour are found
    if (root1 == -1) {
        // make any node (say 1)
        // as the root
        cout << endl
             << "1" << endl;
    }

    // Check if making root1
    // as the root of the
    // tree solves the purpose
    else if (
        checkPossibility(root1, colours)) {

        cout << root1 << endl;
    }

    // check  for root2
    else if (
        checkPossibility(root2, colours)) {

        cout << root2 << endl;
    }

    // otherwise no such root exist
    else {
        cout << "-1" << endl;
    }
}

// Driver code
int32_t main()
{
    // Number of nodes
    int N = 8;

    // add edges
    addedge(1, 2);
    addedge(1, 3);
    addedge(2, 4);
    addedge(2, 7);
    addedge(3, 5);
    addedge(3, 6);
    addedge(6, 8);

    // Node colours
    // 0th node is extra to make
    // the array 1 indexed
    int colours[9] = { 0, 1, 1, 1,
                       1, 1, 2, 1, 3 };

    solve(colours, N);

    return 0;
}

Java

// Java program for the above approach
import java.util.*;

class GFG{

static int NN = (int)(1e5 + 5);

// Vector to store the tree
@SuppressWarnings("unchecked")
static Vector<Integer> []G = new Vector[NN];

// Function to perform dfs
static void dfs(int node, int parent,
                boolean check,
                int current_colour,
                int[] colours)
{

    // Check is assigned to false if either it
    // is already false or the current_colour
    // is not same as the node colour
    check = check &&
           (colours[node] == current_colour);

    // Iterate over the neighbours of node
    for(int a : G[node])
    {

        // If the neighbour is
        // not the parent node
        if (a != parent)
        {

            // Call the function
            // for the neighbour
            dfs(a, node, check,
                current_colour,
                colours);
        }
    }
}

// Function to check whether all the
// nodes in each subtree of the given
// node have same colour
static boolean checkPossibility(int root,
                                int[] colours)
{

    // Initialise the boolean answer
    boolean ans = true;

    // Iterate over the neighbours
    // of selected root
    for(int a : G[root])
    {

        // Initialise the colour
        // for this subtree
        // as the colour of
        // first neighbour
        int current_colour = colours[a];

        // Variable to check
        // condition of same
        // colour for each subtree
        boolean check = true;

        // dfs function call
        dfs(a, root, check,
            current_colour, colours);

        // Check if any one subtree
        // does not have all
        // nodes of same colour
        // then ans will become false
        ans = ans && check;
    }

    // Return the answer
    return ans;
}

// Function to add edges to the tree
static void addedge(int x, int y)
{

    // y is added as a neighbour of x
    G[x].add(y);

    // x is added as a neighbour of y
    G[y].add(x);
}

// Function to find the node
static void solve(int[] colours, int N)
{

    // Initialise root1 as -1
    int root1 = -1;

    // Initialise root2 as -1
    int root2 = -1;

    // Find the first two nodes of
    // different colour which are adjacent
    // to each other
    for(int i = 1; i <= N; i++)
    {
        for(int a : G[i])
        {
            if (colours[a] != colours[i])
            {
                root1 = a;
                root2 = i;
                break;
            }
        }
    }

    // If no two nodes of different
    // colour are found
    if (root1 == -1)
    {

        // Make any node (say 1)
        // as the root
        System.out.println("1" + "\n");
    }

    // Check if making root1
    // as the root of the
    // tree solves the purpose
    else if (checkPossibility(root1, colours))
    {
        System.out.print(root1 + "\n");
    }

    // Check for root2
    else if (checkPossibility(root2, colours))
    {
        System.out.print(root2 + "\n");
    }

    // Otherwise no such root exist
    else
    {
        System.out.print("-1" + "\n");
    }
}

// Driver code
public static void main(String[] args)
{

    // Number of nodes
    int N = 8;

    for(int i = 0; i < G.length; i++)
        G[i] = new Vector<Integer>();

    // Add edges
    addedge(1, 2);
    addedge(1, 3);
    addedge(2, 4);
    addedge(2, 7);
    addedge(3, 5);
    addedge(3, 6);
    addedge(6, 8);

    // Node colours 0th node is extra
    // to make the array 1 indexed
    int colours[] = { 0, 1, 1, 1,
                      1, 1, 2, 1, 3 };

    solve(colours, N);
}
}

// This code is contributed by 29AjayKumar

Python3

# Python3 program for the above approach
NN = 1e5 + 5

# Vector to store tree
G = []
for i in range(int(NN)):
    G.append([])

# Function to perform dfs
def dfs(node, parent, check,
        current_colour, colours):

    # Check is assigned to false if
    # either it is already false or
    # the current_colour is not same
    # as the node colour
    check[0] = check[0] & (colours[node] ==
                           current_colour)

    # Iterate over the neighbours of node
    for a in G[node]:

        # If the neighbour is
        # not the parent node
        if a != parent:

            # Call the function
            # for the neighbour
            dfs(a, node, check,
                current_colour, colours)

# Function to check whether all the
# nodes in each subtree of the given
# node have same colour
def checkPossibility(root, colours):

    # Initialise the boolean answer
    ans = True

    for a in G[root]:

        # Initialise the colour
        # for this subtree
        # as the colour of
        # first neighbour
        current_colour = colours[a]

        # Variable to check
        # condition of same
        # colour for each subtree
        check = [True]

        # dfs function call
        dfs(a, root, check,
            current_colour, colours)

        # Check if any one subtree
        # does not have all
        # nodes of same colour
        # then ans will become false
        ans = ans & check[0]

    # Return the ans
    return ans

# Function to add edges to the tree
def addedge(x, y):

    # y is added as a neighbour of x
    G[x].append(y)

    # x is added as a neighbour of y
    G[y].append(x)

# Function to find the node
def solve(colours, N):

    # Initialise the root1 as -1
    root1 = -1

    # Initialise the root 2 as -1
    root2 = -1

    # Find the first two nodes of
    # different colour which are adjacent
    # to each other
    for i in range(1, N + 1):
        for a in G[i]:
            if colours[a] != colours[i]:
                root1 = a
                root2 = i
                break

    # If no two nodes of different
    # colour are found
    if root1 == -1:

        # make any node (say 1)
        # as the root
        print(1)

    # Check if making root1
    # as the root of the
    # tree solves the purpose
    elif checkPossibility(root1, colours):
        print(root1)

    # Check for root2
    elif checkPossibility(root2, colours):
        print(root2)

    # Otherwise no such root exist
    else:
        print(-1)

# Driver code

# Number of nodes
N = 8

# add edges
addedge(1, 2)
addedge(1, 3)
addedge(2, 4)
addedge(2, 7)
addedge(3, 5)
addedge(3, 6)
addedge(6, 8)

# Node colours
# 0th node is extra to make
# the array 1 indexed
colours = [ 0, 1, 1, 1, 1,
            1, 2, 1, 3 ]

solve(colours, N)

# This code is contributed by Stuti Pathak

C#

// C# program for the above approach
using System;
using System.Collections.Generic;
class GFG{
    static int NN = (int)(1e5 + 5);

    // List to store the tree
    static List<int>[] G = new List<int>[ NN ];

    // Function to perform dfs
    static void dfs(int node, int parent, bool check,
                    int current_colour, int[] colours)
    {

        // Check is assigned to false if either it
        // is already false or the current_colour
        // is not same as the node colour
        check = check && (colours[node] == current_colour);

        // Iterate over the neighbours of node
        foreach(int a in G[node])
        {

            // If the neighbour is
            // not the parent node
            if (a != parent)
            {

                // Call the function
                // for the neighbour
                dfs(a, node, check,
                    current_colour, colours);
            }
        }
    }

    // Function to check whether all the
    // nodes in each subtree of the given
    // node have same colour
    static bool checkPossibility(int root, int[] colours)
    {

        // Initialise the bool answer
        bool ans = true;

        // Iterate over the neighbours
        // of selected root
        foreach(int a in G[root])
        {

            // Initialise the colour
            // for this subtree
            // as the colour of
            // first neighbour
            int current_colour = colours[a];

            // Variable to check
            // condition of same
            // colour for each subtree
            bool check = true;

            // dfs function call
            dfs(a, root, check, current_colour, colours);

            // Check if any one subtree
            // does not have all
            // nodes of same colour
            // then ans will become false
            ans = ans && check;
        }

        // Return the answer
        return ans;
    }

    // Function to add edges to the tree
    static void addedge(int x, int y)
    {

        // y is added as a neighbour of x
        G[x].Add(y);

        // x is added as a neighbour of y
        G[y].Add(x);
    }

    // Function to find the node
    static void solve(int[] colours, int N)
    {

        // Initialise root1 as -1
        int root1 = -1;

        // Initialise root2 as -1
        int root2 = -1;

        // Find the first two nodes of
        // different colour which are adjacent
        // to each other
        for (int i = 1; i <= N; i++)
        {
            foreach(int a in G[i])
            {
                if (colours[a] != colours[i])
                {
                    root1 = a;
                    root2 = i;
                    break;
                }
            }
        }

        // If no two nodes of different
        // colour are found
        if (root1 == -1)
        {

            // Make any node (say 1)
            // as the root
            Console.WriteLine("1" + "\n");
        }

        // Check if making root1
        // as the root of the
        // tree solves the purpose
        else if (checkPossibility(root1, colours))
        {
            Console.Write(root1 + "\n");
        }

        // Check for root2
        else if (checkPossibility(root2, colours))
        {
            Console.Write(root2 + "\n");
        }

        // Otherwise no such root exist
        else
        {
            Console.Write("-1" + "\n");
        }
    }

    // Driver code
    public static void Main(String[] args)
    {

        // Number of nodes
        int N = 8;

        for (int i = 0; i < G.Length; i++)
            G[i] = new List<int>();

        // Add edges
        addedge(1, 2);
        addedge(1, 3);
        addedge(2, 4);
        addedge(2, 7);
        addedge(3, 5);
        addedge(3, 6);
        addedge(6, 8);

        // Node colours 0th node is extra
        // to make the array 1 indexed
        int[] colours = {0, 1, 1, 1, 1, 1, 2, 1, 3};
        solve(colours, N);
    }
}

// This code is contributed by Rajput-Ji

Javascript

<script>

// Javascript program for the above approach

var NN = 100005;

// List to store the tree
var G = Array.from(Array(NN), ()=>Array());

// Function to perform dfs
function dfs(node, parent, check, current_colour, colours)
{
    // Check is assigned to false if either it
    // is already false or the current_colour
    // is not same as the node colour
    check = check && (colours[node] == current_colour);
    // Iterate over the neighbours of node
    for(var a of G[node])
    {
        // If the neighbour is
        // not the parent node
        if (a != parent)
        {
            // Call the function
            // for the neighbour
            dfs(a, node, check,
                current_colour, colours);
        }
    }
}
// Function to check whether all the
// nodes in each subtree of the given
// node have same colour
function checkPossibility(root, colours)
{
    // Initialise the bool answer
    var ans = true;
    // Iterate over the neighbours
    // of selected root
    for(var a of G[root])
    {
        // Initialise the colour
        // for this subtree
        // as the colour of
        // first neighbour
        var current_colour = colours[a];
        // Variable to check
        // condition of same
        // colour for each subtree
        var check = true;
        // dfs function call
        dfs(a, root, check, current_colour, colours);
        // Check if any one subtree
        // does not have all
        // nodes of same colour
        // then ans will become false
        ans = ans && check;
    }
    // Return the answer
    return ans;
}
// Function to add edges to the tree
function addedge(x, y)
{
    // y is added as a neighbour of x
    G[x].push(y);
    // x is added as a neighbour of y
    G[y].push(x);
}
// Function to find the node
function solve(colours, N)
{
    // Initialise root1 as -1
    var root1 = -1;
    // Initialise root2 as -1
    var root2 = -1;
    // Find the first two nodes of
    // different colour which are adjacent
    // to each other
    for (var i = 1; i <= N; i++)
    {
        for(var a of G[i])
        {
            if (colours[a] != colours[i])
            {
                root1 = a;
                root2 = i;
                break;
            }
        }
    }
    // If no two nodes of different
    // colour are found
    if (root1 == -1)
    {
        // Make any node (say 1)
        // as the root
        document.write("1" + "<br>");
    }
    // Check if making root1
    // as the root of the
    // tree solves the purpose
    else if (checkPossibility(root1, colours))
    {
        document.write(root1 + "<br>");
    }
    // Check for root2
    else if (checkPossibility(root2, colours))
    {
        document.write(root2 + "<br>");
    }
    // Otherwise no such root exist
    else
    {
        document.write("-1" + "<br>");
    }
}
// Driver code
// Number of nodes
var N = 8;

// Add edges
addedge(1, 2);
addedge(1, 3);
addedge(2, 4);
addedge(2, 7);
addedge(3, 5);
addedge(3, 6);
addedge(6, 8);
// Node colours 0th node is extra
// to make the array 1 indexed
var colours = [0, 1, 1, 1, 1, 1, 2, 1, 3];
solve(colours, N);


</script>

Output:

6

Time Complexity: O(N*N) where N is the number of nodes in the tree. 
Auxiliary Space: O(N)