Count of nodes in a given N-ary tree having distance to all leaf nodes equal in their subtree - GeeksforGeeks

11 Jan, 2022
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 imageOutput: 4Explanation: 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 belowOutput: 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 parentIf every branch has equal height then increment the count by 1 and return height + 1 to the parentOr else return -1 to the parent indicating that the branches have unequal height
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++
Java
Python3
C#
Javascript
// 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 treeint 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 heightint 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 codeint 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 approachclass Node:    def __init__(self, val):        self.val = val        self.children = [] # Post-order traversal to find# depth of all branches of every# node of the treedef 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 heightdef 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 treeseven = 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 resultprint(equalHeightBranches(seven)) # This code is contributed by rj13to.
// C# implementation for the above approachusing 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 approachclass 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 heightfunction 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 treefunction 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 treelet 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 resultconsole.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 
gfgking
lokeshpotta20
simmytarika5
simranarora5sos
SHUBHAMSINGH10
rj13to
n-ary-tree
tree-traversal
Tree
Tree
Writing code in comment?
Please use ide.geeksforgeeks.org,
generate link and share the link here.
Comments
Old Comments
Introduction to Tree Data Structure
DFS traversal of a tree using recursion
Top 50 Tree Coding Problems for Interviews
Check if a given graph is tree or not
Difference between Min Heap and Max Heap
Iterative Preorder Traversal
Find the node with minimum value in a Binary Search Tree
Iterative Postorder Traversal | Set 2 (Using One Stack)
Real-time application of Data Structures
Overview of Data Structures | Set 2 (Binary Tree, BST, Heap and Hash)