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 / 4Output: 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 7Output: 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> |
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.