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

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

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

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

Python3

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

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

Javascript

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

Output

4

Time Complexity: O(N)
Auxiliary Space: O(H), where H is the height of the tree