Count pairs of leaf nodes in a Binary Tree which are at most K distance apart

Given a binary tree and an integer K, the task is to count the possible pairs of leaf nodes from the given binary tree such that the distance between them is at most K.

Examples:

Input: K = 3 
 

      1
     / \
    2   3
   /
  4


Output:
Explanation: 
The leaf nodes of the tree are 3 and 4 
And the shortest distance between them is 3. 
This is the only valid pair.
Input: K = 3 
 

       1
      / \
     2   3
    / \ / \
   4  5 6  7


Output: 2
 



Approach: The idea is to use post order traversal with an array of size K + 1 to keep track of the number of nodes with a particular distance. Below are the steps:
 

  • Initialize an array arr[] of size K + 1, where arr[i] denotes the number of leaf nodes at distance i from the current node.
  • Then, recursively update the above array for the left and the right subtree of every node in an array left[] and right[] respectively.
  • After the above step, for each node the distance between left[] and right[] at the corresponding index will give the distance between the leftmost and rightmost leaf node. Update it in the array res[] as: 

res[i + 1] = left[i] * right[i] 
 

  • Now for all possible pairs (l, r) in the arrays left[] and right[], if the sum of the distance between them is at most K, then increment the count of pairs as: 

count += left[l] * right[r] 
 

Below is the implementation of the above approach:  

C++14

filter_none

edit
close

play_arrow

link
brightness_4
code

// C++14 implementation of the
// above approach
#include <bits/stdc++.h>
using namespace std;
 
// Structure of a Node
struct Node
{
    int data;
    Node* left, *right;
 
    // Constructor of the class
    Node(int item)
    {
        data = item;
        left = right = NULL;
    }
};
 
// Stores the count of required pairs
int result = 0;
 
// Function to perform dfs to find pair of
// leaf nodes at most K distance apart
vector<int> dfs(Node* root, int distance)
{
     
    // Return empty array if node is NULL
    if (root == NULL)
    {
        vector<int> res(distance + 1, 0);
        return res;
    }
 
    // If node is a leaf node and return res
    if (root->left == NULL &&
       root->right == NULL)
    {
        vector<int> res(distance + 1, 0);
        res[1]++;
        return res;
    }
 
    // Traverse to the left
    vector<int> left = dfs(root->left,
                           distance);
 
    // Traverse to the right
    vector<int> right = dfs(root->right,
                            distance);
 
    vector<int> res(distance + 1, 0);
 
    // Update the distance between left
    // and right leaf node
    for(int i = res.size() - 2;
            i >= 1; i--)
        res[i + 1] = left[i]+ right[i];
 
    // Count all pair of leaf nodes
    // which are at most K distance apart
    for(int l = 1;l < left.size(); l++)
    {
        for(int r = 0;r < right.size(); r++)
        {
            if (l + r <= distance)
            {
                result += left[l] * right[r];
            }
        }
    }
 
    // Return res to parent node
    return res;
}
 
// Driver Code
int main()
{
     
    /*
         1
        / \
       2   3
      /
    4
    */
 
    Node* root = new Node(1);
    root->left = new Node(2);
    root->right = new Node(3);
    root->left->left = new Node(4);
 
    // Given distance K
    int K = 3;
 
    // Function call
    dfs(root, K);
 
    cout << result;
}
 
// This code is contributed by mohit kumar 29

chevron_right


Java

filter_none

edit
close

play_arrow

link
brightness_4
code

// Java implementation of the
// above approach
 
// Structure of a Node
class Node {
    int data;
    Node left, right;
 
    // Constructor of the class
    public Node(int item)
    {
        data = item;
        left = right = null;
    }
}
 
public class GFG {
    Node root;
 
    // Stores the count of required pairs
    static int result;
 
    // Function to perform dfs to find pair of
    // leaf nodes at most K distance apart
    static int[] dfs(Node root, int distance)
    {
        // Return empty array if node is null
        if (root == null)
            return new int[distance + 1];
 
        // If node is a leaf node and return res
        if (root.left == null
            && root.right == null) {
            int res[] = new int[distance + 1];
            res[1]++;
            return res;
        }
 
        // Traverse to the left
        int[] left = dfs(root.left,
                         distance);
 
        // Traverse to the right
        int[] right = dfs(root.right,
                          distance);
 
        int res[] = new int[distance + 1];
 
        // Update the distance between left
        // and right leaf node
        for (int i = res.length - 2;
             i >= 1; i--) {
            res[i + 1] = left[i]
                         + right[i];
        }
 
        // Count all pair of leaf nodes
        // which are at most K distance apart
        for (int l = 1;
             l < left.length; l++) {
 
            for (int r = 0;
                 r < right.length; r++) {
 
                if (l + r <= distance) {
                    result += left[l]
                              * right[r];
                }
            }
        }
 
        // Return res to parent node
        return res;
    }
 
    // Driver Code
    public static void main(String[] args)
    {
        GFG tree = new GFG();
 
        /*
                1
               /  \
             2     3
            
           4   
       */
 
        tree.root = new Node(1);
        tree.root.left = new Node(2);
        tree.root.right = new Node(3);
 
        tree.root.left.left = new Node(4);
        result = 0;
 
        // Given distance K
        int K = 3;
 
        // Function Call
        dfs(tree.root, K);
 
        System.out.println(result);
    }
}

chevron_right


C#

filter_none

edit
close

play_arrow

link
brightness_4
code

// C# implementation of the
// above approach
using System;
 
// Structure of a Node
class Node{
   
  public int data;
  public Node left, right;
 
  // Constructor of the class
  public Node(int item)
  {
    data = item;
    left = right = null;
  }
}
 
class GFG{
 
Node root;
 
// Stores the count of
// required pairs
static int result;
 
// Function to perform
// dfs to find pair of
// leaf nodes at most K
// distance apart
static int[] dfs(Node root,
                 int distance)
{
  int []res;
  // Return empty array if
  //  node is null
  if (root == null)
    return new int[distance + 1];
 
  // If node is a leaf node
  //and return res
  if (root.left == null &&
      root.right == null)
  {
    res = new int[distance + 1];
    res[1]++;
    return res;
  }
 
  // Traverse to the left
  int[] left = dfs(root.left,
                   distance);
 
  // Traverse to the right
  int[] right = dfs(root.right,
                    distance);
 
  res = new int[distance + 1];
 
  // Update the distance between left
  // and right leaf node
  for (int i = res.Length - 2;
           i >= 1; i--)
  {
    res[i + 1] = left[i] + right[i];
  }
 
  // Count all pair of leaf nodes
  // which are at most K distance apart
  for (int l = 1; l < left.Length; l++)
  {
    for (int r = 0; r < right.Length; r++)
    {
      if (l + r <= distance)
      {
        result += left[l] * right[r];
      }
    }
  }
 
  // Return res to parent node
  return res;
}
 
    // Driver Code
public static void Main(String[] args)
{
  GFG tree = new GFG();
 
  /*
                1
               /  \
             2     3
            
           4   
       */
 
  tree.root = new Node(1);
  tree.root.left = new Node(2);
  tree.root.right = new Node(3);
 
  tree.root.left.left = new Node(4);
  result = 0;
 
  // Given distance K
  int K = 3;
 
  // Function Call
  dfs(tree.root, K);
 
  Console.WriteLine(result);
}
}
 
// This code is contributed by Rajput-Ji

chevron_right


Output: 

1


 

Time Complexity: O(N + E), where N is the number of nodes in the Binary tree, and E is the number of Edges.Auxiliary Space: O(H), where H is the height of the Binary tree.

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 : mohit kumar 29, Rajput-Ji