Given an N-ary tree root, the task is to find the number of non-leaf nodes in the tree such that all the leaf nodes in the subtree of the current node are at an equal distance from the current node.
Example:
Input: Tree in the below image
Output: 4
Explanation: The nodes {10, 3, 2, 4} have the distance between them and all the leaf nodes in their subtree respectively as equal.
Input: Tree in the image below
Output: 3
Approach: The given problem can be solved by using the post-order traversal. The idea is to check if the number of nodes from the current node to all its leaf nodes is the same. Below steps can be followed to solve the problem:
- Apply post-order traversal on the N-ary tree:
- If the root has no children then return 1 to the parent
- If every branch has equal height then increment the count by 1 and return height + 1 to the parent
- Or else return -1 to the parent indicating that the branches have unequal height
- Return the count as the answer
// C++ code for the above approach #include <bits/stdc++.h> using namespace std;
class Node {
public :
vector<Node*> children;
int val;
// constructor
Node( int v)
{
val = v;
children = {};
}
}; // Post-order traversal to find // depth of all branches of every // node of the tree int postOrder(Node* root, int count[])
{ // If root is a leaf node
// then return 1
if (root->children.size() == 0)
return 1;
// Initialize a variable height
// calculate longest increasing path
int height = 0;
// Use recursion on all child nodes
for (Node* child : root->children)
{
// Get the height of the branch
int h = postOrder(child, count);
// Initialize height of first
// explored branch
if (height == 0)
height = h;
// If branches are unbalanced
// then store -1 in height
else if (h == -1 || height != h)
height = -1;
}
// Increment the value of count
// If height is not -1
if (height != -1)
count[0]++;
// Return the height of branches
// including the root if height is
// not -1 or else return -1
return height != -1 ? height + 1 : -1;
} // Function to find the number of nodes // in the N-ary tree with their branches // having equal height int equalHeightBranches(Node* root)
{ // Base case
if (root == NULL)
return 0;
// Initialize a variable count
// to store the answer
int count[1] = { 0 };
// Apply post order traversal
// on the tree
postOrder(root, count);
// Return the answer
return count[0];
} // Driver code int main()
{ // Initialize the tree
Node* seven = new Node(7);
Node* seven2 = new Node(7);
Node* five = new Node(5);
Node* four = new Node(4);
Node* nine = new Node(9);
Node* one = new Node(1);
Node* two = new Node(2);
Node* six = new Node(6);
Node* eight = new Node(8);
Node* ten = new Node(10);
Node* three = new Node(3);
Node* mfour = new Node(-4);
Node* mtwo = new Node(-2);
Node* zero = new Node(0);
three->children.push_back(mfour);
three->children.push_back(mtwo);
three->children.push_back(zero);
ten->children.push_back(three);
two->children.push_back(six);
two->children.push_back(seven2);
four->children.push_back(nine);
four->children.push_back(one);
four->children.push_back(five);
seven->children.push_back(ten);
seven->children.push_back(two);
seven->children.push_back(eight);
seven->children.push_back(four);
// Call the function
// and print the result
cout << (equalHeightBranches(seven));
} // This code is contributed by Potta Lokesh |
// Java implementation for the above approach import java.io.*;
import java.util.*;
class GFG {
// Function to find the number of nodes
// in the N-ary tree with their branches
// having equal height
public static int equalHeightBranches(Node root)
{
// Base case
if (root == null )
return 0 ;
// Initialize a variable count
// to store the answer
int [] count = new int [ 1 ];
// Apply post order traversal
// on the tree
postOrder(root, count);
// Return the answer
return count[ 0 ];
}
// Post-order traversal to find
// depth of all branches of every
// node of the tree
public static int postOrder(
Node root, int [] count)
{
// If root is a leaf node
// then return 1
if (root.children.size() == 0 )
return 1 ;
// Initialize a variable height
// calculate longest increasing path
int height = 0 ;
// Use recursion on all child nodes
for (Node child : root.children) {
// Get the height of the branch
int h = postOrder(child, count);
// Initialize height of first
// explored branch
if (height == 0 )
height = h;
// If branches are unbalanced
// then store -1 in height
else if (h == - 1 || height != h)
height = - 1 ;
}
// Increment the value of count
// If height is not -1
if (height != - 1 )
count[ 0 ]++;
// Return the height of branches
// including the root if height is
// not -1 or else return -1
return height != - 1 ? height + 1 : - 1 ;
}
// Driver code
public static void main(String[] args)
{
// Initialize the tree
Node seven = new Node( 7 );
Node seven2 = new Node( 7 );
Node five = new Node( 5 );
Node four = new Node( 4 );
Node nine = new Node( 9 );
Node one = new Node( 1 );
Node two = new Node( 2 );
Node six = new Node( 6 );
Node eight = new Node( 8 );
Node ten = new Node( 10 );
Node three = new Node( 3 );
Node mfour = new Node(- 4 );
Node mtwo = new Node(- 2 );
Node zero = new Node( 0 );
three.children.add(mfour);
three.children.add(mtwo);
three.children.add(zero);
ten.children.add(three);
two.children.add(six);
two.children.add(seven2);
four.children.add(nine);
four.children.add(one);
four.children.add(five);
seven.children.add(ten);
seven.children.add(two);
seven.children.add(eight);
seven.children.add(four);
// Call the function
// and print the result
System.out.println(
equalHeightBranches(seven));
}
static class Node {
List<Node> children;
int val;
// constructor
public Node( int val)
{
this .val = val;
children = new ArrayList<>();
}
}
} |
# Python code for the above approach class Node:
def __init__( self , val):
self .val = val
self .children = []
# Post-order traversal to find # depth of all branches of every # node of the tree def postOrder(root, count):
# If root is a leaf node
# then return 1
if ( len (root.children) = = 0 ):
return 1
# Initialize a variable height
# calculate longest increasing path
height = 0
# Use recursion on all child nodes
for child in root.children:
# Get the height of the branch
h = postOrder(child, count)
# Initialize height of first
# explored branch
if (height = = 0 ):
height = h
# If branches are unbalanced
# then store -1 in height
elif (h = = - 1 or height ! = h):
height = - 1
# Increment the value of count
# If height is not -1
if (height ! = - 1 ):
count[ 0 ] + = 1
# Return the height of branches
# including the root if height is
# not -1 or else return -1
if (height ! = - 1 ):
return height + 1
else :
return - 1
# Function to find the number of nodes # in the N-ary tree with their branches # having equal height def equalHeightBranches(root):
# Base case
if (root = = None ):
return 0
# Initialize a variable count
# to store the answer
count = [ 0 ]
# Apply post order traversal
# on the tree
postOrder(root, count)
# Return the answer
return count[ 0 ]
# Driver code # Initialize the tree seven = Node( 7 )
seven2 = Node( 7 )
five = Node( 5 )
four = Node( 4 )
nine = Node( 9 )
one = Node( 1 )
two = Node( 2 )
six = Node( 6 )
eight = Node( 8 )
ten = Node( 10 )
three = Node( 3 )
mfour = Node( - 4 )
mtwo = Node( - 2 )
zero = Node( 0 )
three.children.append(mfour) three.children.append(mtwo) three.children.append(zero) ten.children.append(three) two.children.append(six) two.children.append(seven2) four.children.append(nine) four.children.append(one) four.children.append(five) seven.children.append(ten) seven.children.append(two) seven.children.append(eight) seven.children.append(four) # Call the function # and print the result print (equalHeightBranches(seven))
# This code is contributed by rj13to. |
// C# implementation for the above approach using System;
using System.Collections.Generic;
public class Node
{ public int val;
public List<Node> children;
// Constructor to create a Node
public Node( int vall)
{
val = vall;
children = new List<Node>();
}
} class GFG {
// Function to find the number of nodes
// in the N-ary tree with their branches
// having equal height
public static int equalHeightBranches(Node root)
{
// Base case
if (root == null )
return 0;
// Initialize a variable count
// to store the answer
int [] count = new int [1];
// Apply post order traversal
// on the tree
postOrder(root, count);
// Return the answer
return count[0];
}
// Post-order traversal to find
// depth of all branches of every
// node of the tree
public static int postOrder(
Node root, int [] count)
{
// If root is a leaf node
// then return 1
if (root.children.Count == 0)
return 1;
// Initialize a variable height
// calculate longest increasing path
int height = 0;
// Use recursion on all child nodes
foreach (Node child in root.children) {
// Get the height of the branch
int h = postOrder(child, count);
// Initialize height of first
// explored branch
if (height == 0)
height = h;
// If branches are unbalanced
// then store -1 in height
else if (h == -1 || height != h)
height = -1;
}
// Increment the value of count
// If height is not -1
if (height != -1)
count[0]++;
// Return the height of branches
// including the root if height is
// not -1 or else return -1
return height != -1 ? height + 1 : -1;
}
// Driver code
public static void Main()
{
// Initialize the tree
Node seven = new Node(7);
Node seven2 = new Node(7);
Node five = new Node(5);
Node four = new Node(4);
Node nine = new Node(9);
Node one = new Node(1);
Node two = new Node(2);
Node six = new Node(6);
Node eight = new Node(8);
Node ten = new Node(10);
Node three = new Node(3);
Node mfour = new Node(-4);
Node mtwo = new Node(-2);
Node zero = new Node(0);
three.children.Add(mfour);
three.children.Add(mtwo);
three.children.Add(zero);
ten.children.Add(three);
two.children.Add(six);
two.children.Add(seven2);
four.children.Add(nine);
four.children.Add(one);
four.children.Add(five);
seven.children.Add(ten);
seven.children.Add(two);
seven.children.Add(eight);
seven.children.Add(four);
// Call the function
// and print the result
Console.WriteLine(
equalHeightBranches(seven));
}
} // This code is contributed // by Shubham Singh |
<script> // Javascript implementation for the above approach class Node { // constructor
constructor(val) {
this .val = val;
this .children = [];
}
} // Function to find the number of nodes // in the N-ary tree with their branches // having equal height function equalHeightBranches(root) {
// Base case
if (root == null )
return 0;
// Initialize a variable count
// to store the answer
let count = [0];
// Apply post order traversal
// on the tree
postOrder(root, count);
// Return the answer
return count[0];
} // Post-order traversal to find // depth of all branches of every // let of the tree function postOrder(root, count) {
// If root is a leaf node
// then return 1
if (root.children.length == 0)
return 1;
// Initialize a variable height
// calculate longest increasing path
let height = 0;
// Use recursion on all child nodes
for (child of root.children) {
// Get the height of the branch
let h = postOrder(child, count);
// Initialize height of first
// explored branch
if (height == 0)
height = h;
// If branches are unbalanced
// then store -1 in height
else if (h == -1 || height != h)
height = -1;
}
// Increment the value of count
// If height is not -1
if (height != -1)
count[0]++;
// Return the height of branches
// including the root if height is
// not -1 or else return -1
return height != -1 ? height + 1 : -1;
} // Driver code // Initialize the tree let seven = new Node(7);
let seven2 = new Node(7);
let five = new Node(5);
let four = new Node(4);
let nine = new Node(9);
let one = new Node(1);
let two = new Node(2);
let six = new Node(6);
let eight = new Node(8);
let ten = new Node(10);
let three = new Node(3);
let mfour = new Node(-4);
let mtwo = new Node(-2);
let zero = new Node(0);
three.children.push(mfour); three.children.push(mtwo); three.children.push(zero); ten.children.push(three); two.children.push(six); two.children.push(seven2); four.children.push(nine); four.children.push(one); four.children.push(five); seven.children.push(ten); seven.children.push(two); seven.children.push(eight); seven.children.push(four); // Call the function // and print the result console.log(equalHeightBranches(seven)); // This code is contributed by gfgking. </script> |
4
Time Complexity: O(N)
Auxiliary Space: O(H), where H is the height of the tree