Count pairs of leaf nodes in a Binary Tree which are at most K distance apart
Last Updated :
21 Jun, 2021
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++
#include <bits/stdc++.h>
using namespace std;
struct Node
{
int data;
Node* left, *right;
Node( int item)
{
data = item;
left = right = NULL;
}
};
int result = 0;
vector< int > dfs(Node* root, int distance)
{
if (root == NULL)
{
vector< int > res(distance + 1, 0);
return res;
}
if (root->left == NULL &&
root->right == NULL)
{
vector< int > res(distance + 1, 0);
res[1]++;
return res;
}
vector< int > left = dfs(root->left,
distance);
vector< int > right = dfs(root->right,
distance);
vector< int > res(distance + 1, 0);
for ( int i = res.size() - 2;
i >= 1; i--)
res[i + 1] = left[i]+ right[i];
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;
}
int main()
{
Node* root = new Node(1);
root->left = new Node(2);
root->right = new Node(3);
root->left->left = new Node(4);
int K = 3;
dfs(root, K);
cout << result;
}
|
Java
class Node {
int data;
Node left, right;
public Node( int item)
{
data = item;
left = right = null ;
}
}
public class GFG {
Node root;
static int result;
static int [] dfs(Node root, int distance)
{
if (root == null )
return new int [distance + 1 ];
if (root.left == null
&& root.right == null ) {
int res[] = new int [distance + 1 ];
res[ 1 ]++;
return res;
}
int [] left = dfs(root.left,
distance);
int [] right = dfs(root.right,
distance);
int res[] = new int [distance + 1 ];
for ( int i = res.length - 2 ;
i >= 1 ; i--) {
res[i + 1 ] = left[i]
+ right[i];
}
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;
}
public static void main(String[] args)
{
GFG tree = new GFG();
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 ;
int K = 3 ;
dfs(tree.root, K);
System.out.println(result);
}
}
|
Python3
class newNode:
def __init__( self , item):
self .data = item
self .left = None
self .right = None
result = 0
def dfs(root, distance):
global result
if (root = = None ):
res = [ 0 for i in range (distance + 1 )]
return res
if (root.left = = None and root.right = = None ):
res = [ 0 for i in range (distance + 1 )]
res[ 1 ] + = 1
return res
left = dfs(root.left, distance)
right = dfs(root.right, distance)
res = [ 0 for i in range (distance + 1 )]
i = len (res) - 2
while (i > = 1 ):
res[i + 1 ] = left[i] + right[i]
i - = 1
for l in range ( 1 , len (left)):
for r in range ( len (right)):
if (l + r < = distance):
result + = left[l] * right[r]
return res
if __name__ = = '__main__' :
root = newNode( 1 )
root.left = newNode( 2 )
root.right = newNode( 3 )
root.left.left = newNode( 4 )
K = 3
dfs(root, K)
print (result)
|
C#
using System;
class Node{
public int data;
public Node left, right;
public Node( int item)
{
data = item;
left = right = null ;
}
}
class GFG{
Node root;
static int result;
static int [] dfs(Node root,
int distance)
{
int []res;
if (root == null )
return new int [distance + 1];
if (root.left == null &&
root.right == null )
{
res = new int [distance + 1];
res[1]++;
return res;
}
int [] left = dfs(root.left,
distance);
int [] right = dfs(root.right,
distance);
res = new int [distance + 1];
for ( int i = res.Length - 2;
i >= 1; i--)
{
res[i + 1] = left[i] + right[i];
}
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;
}
public static void Main(String[] args)
{
GFG tree = new GFG();
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;
int K = 3;
dfs(tree.root, K);
Console.WriteLine(result);
}
}
|
Javascript
<script>
class Node
{
constructor(item) {
this .left = null ;
this .right = null ;
this .data = item;
}
}
let root;
class GFG{
}
let result;
function dfs(root, distance)
{
if (root == null )
{
let arr = new Array(distance + 1);
arr.fill(0);
return arr;
}
if (root.left == null
&& root.right == null ) {
let res = new Array(distance + 1);
res.fill(0);
res[1]++;
return res;
}
let left = dfs(root.left, distance);
let right = dfs(root.right, distance);
let res = new Array(distance + 1);
res.fill(0);
for (let i = res.length - 2; i >= 1; i--) {
res[i + 1] = left[i] + right[i];
}
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;
}
let tree = new GFG();
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;
let K = 3;
dfs(tree.root, K);
document.write(result);
</script>
|
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.
Share your thoughts in the comments
Please Login to comment...