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

Given a tree a with N nodes numbered from 1 to N and N – 1 edge and an array colours[] where colours[i] denote the colour of ith node. The task is to find a node such that each neighbouring tree connected to this node consists of the same coloured nodes. If no such node exist 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)}
 

Visualising 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 colour and the tree rooted at 8 has only one node. Therefore node 6 is one such node.

Input: N = 4, colours[] = {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 colour, then any node can be the root. Otherwise, pick any two nodes which are adjacent to each other and have different colours 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 colour 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 colour and any node can then be taken as root.
  2. Check if the of all nodes of each subtree have the same colour 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++

filter_none

edit
close

play_arrow

link
brightness_4
code

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

chevron_right


Python3

filter_none

edit
close

play_arrow

link
brightness_4
code

# 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

chevron_right


Output: 

6

 

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

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.




My Personal Notes arrow_drop_up

Check out this Author's contributed articles.

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



Improved By : stutipathak31jan

Article Tags :
Practice Tags :


Be the First to upvote.


Please write to us at contribute@geeksforgeeks.org to report any issue with the above content.