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

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

Java

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

Python3

# Python3 implementation of the
# above approach

# Structure of a Node
class newNode:

    def __init__(self, item):

        self.data =  item
        self.left = None
        self.right = None

# Stores the count of required pairs
result = 0

# Function to perform dfs to find pair of
# leaf nodes at most K distance apart
def dfs(root, distance):

    global result

    # Return empty array if node is None
    if (root == None):
        res = [0 for i in range(distance + 1)]
        return res

    # If node is a leaf node and return res
    if (root.left == None and root.right == None):
        res = [0 for i in range(distance + 1)]
        res[1] += 1
        return res

    # Traverse to the left
    left = dfs(root.left, distance)

    # Traverse to the right
    right = dfs(root.right, distance)

    res = [0 for i in range(distance + 1)]

    # Update the distance between left
    # and right leaf node
    i = len(res) - 2

    while(i >= 1):
        res[i + 1] = left[i] + right[i]
        i -= 1

    # Count all pair of leaf nodes
    # which are at most K distance apart
    for l in range(1, len(left)):
        for r in range(len(right)):
            if (l + r <= distance):
                result += left[l] * right[r]

    # Return res to parent node
    return res

# Driver Code
if __name__ == '__main__':

    '''
         1
        / \
       2   3
      /
    4
    '''
    root = newNode(1)
    root.left = newNode(2)
    root.right = newNode(3)
    root.left.left = newNode(4)

    # Given distance K
    K = 3

    # Function call
    dfs(root, K)

    print(result)

# This code is contributed by ipg2016107

C#

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

Javascript

<script>

    // JavaScript implementation of the above approach

    // Structure of a Node
    class Node
    {
        constructor(item) {
           this.left = null;
           this.right = null;
           this.data = item;
        }
    }

    let root;

    class GFG{
    }

    // Stores the count of required pairs
    let result;

    // Function to perform dfs to find pair of
    // leaf nodes at most K distance apart
    function dfs(root, distance)
    {
        // Return empty array if node is null
        if (root == null)
        {
            let arr = new Array(distance + 1);
            arr.fill(0);
            return arr;
        }

        // If node is a leaf node and return res
        if (root.left == null
            && root.right == null) {
            let res = new Array(distance + 1);
            res.fill(0);
            res[1]++;
            return res;
        }

        // Traverse to the left
        let left = dfs(root.left, distance);

        // Traverse to the right
        let right = dfs(root.right, distance);

        let res = new Array(distance + 1);
        res.fill(0);

        // Update the distance between left
        // and right leaf node
        for (let 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 (let l = 1; l < left.length; l++) {

            for (let r = 0; r < right.length; r++) {

                if (l + r <= distance) {
                    result += left[l] * right[r];
                }
            }
        }

        // Return res to parent node
        return res;
    }

    let 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
    let K = 3;

    // Function Call
    dfs(tree.root, K);

    document.write(result);

</script>

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.