This blog will discuss the problem "Find the number of pair of ideal nodes in a given tree." In this problem, we will be given an integer 'K' and a tree with 'N' nodes, each node having a value between 1 and N associated with it. We have to find the number of pairs of ideal nodes in the given tree.

A pair of nodes (u,v) is called an ideal pair if it satisfies the following conditions:

u is an ancestor of v

abs(u-v) <= K

Let's take an example to understand the problem:

Assume we have a given binary tree as shown below:

And given K = 3

(1, 2), (1, 3), (1, 4), (3, 5), and (3, 6) are such pairs of (u, v) where u is an ancestor of v and abs(u-v) <= 3

So, here the number of pair of ideal nodes in a given tree is 5

Now that we understand the problem let's discuss the approach to solve it in the next section.

Solution Approach

This section will discuss the approach to find the number of pair of ideal nodes in a given tree. For each node, we need to store their ancestors and then check the value of each ancestor, which satisfies the criteria that the absolute difference of the values of the node and the ancestor is less than or equal to K. In simple words, for each node, we need to find the number of ancestor nodes having a value in a range of [node - K, node + K], where 'node' is the value of the current node. So, we need to perform three major operations -

While traversing the tree, mark a node as ancestor of its descendant nodes.

After traversing all the descendants of a node, remove a node from ancestor nodes.

Find the number of nodes having which are currently marked as the ancestor and have value in a range [node - K, node + K]

We know that the Binary indexed tree, also known as Fenwick tree, is an appropriate data structure using which we can perform the above three operations efficiently in the following way:

To mark a node as an ancestor, we have to insert that node in the binary indexed tree by storing 1 for that node in the array.

To mark a node as not an ancestor, remove the node from the Binary indexed tree by storing 0 for that node in the array.

To find the number of nodes having which are currently marked as an ancestor and have value in a range [node - K, node + K], calculate the sum of the range [node - K, node + K] in the Binary indexed tree.

Hence, this approach can be used to solve the given problem using a binary indexed tree.

Algorithm -

Step 1. First, write the main function. Inside the main function, create an adjacency list to store the edges of the given tree. Also, create a boolean vector "isRoot," which will store 'true' at ith index if node 'i' is the root of the given tree, else store false. Initially store 'true' at all the indices in the vector, then after adding an edge "u->v" in the adjacency list, update 'isRoot[v] = false' as 'v' has its parent 'u,' so 'v' can't be a root.

Step 2. Next, create a vector 'bit' inside the main function for storing indexes of the tree to form the Binary indexed tree. Then create a global variable "idealPairsCount" to store the number of pair of ideal nodes in a given tree, initialize it to zero inside the main function and call the function to find the number of pair of ideal nodes in a given tree.

Step 3. Create a global variable function "findIdealPairs()," which will update the variable 'idealPairsCount" as the count of the total number of pairs of ideal nodes in the given tree. It takes the following input:

Node - The current node being considered (start from the root node)

adj - The adjacency list representation of the given tree

K - Given value of K

bit - vector to store the indexing of the tree

N - Number of nodes in the given tree

Step 4. Inside the function "findIdealNodes()," first call the function to find the number of ancestors of the current node which have values in the range [node-K, node+K], where 'node' is the value of the current node and add it to the variable "idealPairsCount ."Then update the index of the current node to mark it as an ancestor and then call the function "findIdealNodes()" recursively for all its children nodes. After calculating for all the descendants in a DFS Algorithm manner, update its index to remove it from being an ancestor of other nodes.

Step 5. Finally, create two helper functions, "countNodes" to count the number of ancestors having values in a given range and "updateNode()" to update the index of the nodes to mark them as an ancestor or not an ancestor.

C++ code:

// C++ program to to find the number of pair of ideal nodes in a given tree
#include <bits/stdc++.h>
using namespace std;
// Variable to store the number of pair of ideal nodes
long long idealPairsCount;
// Function to return the count of nodes having value between i and j both inclusive
long long countNodes(int i, int j, vector<long long> bit)
{
long long count = 0ll;
while (j > 0)
{
count += bit[j];
j -= (j & (j * -1));
}
i--;
while (i > 0)
{
count -= bit[i];
i -= (i & (i * -1));
}
return count;
}
void updateNode(int i, long long diff, vector<long long>& bit, int n)
{
while (i <= n)
{
bit[i] += diff;
i += i & -i;
}
}
// Function to find the number of pair of ideal nodes in a given tree
void findIdealpairs(int node, vector<int> adj[], int K, vector<long long>& bit, int N)
{
idealPairsCount += countNodes(max(1, node - K),min(N, node + K), bit);
// Update the node to mark it as ancestor of upcoming nodes
updateNode(node, 1, bit, N);
for (int i = 0; i < adj[node].size(); i++)
{
findIdealpairs(adj[node][i], adj, K, bit, N);
}
// After completing all the descendants of current node, update it to mark as not an ancestor
updateNode(node, -1, bit, N);
}
// Main Function
int main()
{
// Number of nodes in the tree
int N = 6;
// Adjacency list to store the edges of the tree
vector<int> adj[N+1];
// Boolean array to mark the root of the tree node
vector<bool> isRoot(N+1, true);
// Inserting the edges in the adjacency list and marking the isRoot[y] as false if there is an edge x->y
adj[1].push_back(2);
isRoot[2] = false;
adj[1].push_back(3);
isRoot[3] = false;
adj[1].push_back(4);
isRoot[4] = false;
adj[3].push_back(5);
isRoot[5] = false;
adj[3].push_back(6);
isRoot[6] = false;
// Given value of K
int K = 3;
// Find the root node of the tree
int root = -1;
for(int i=1; i<=N; i++)
{
if(isRoot[i]==true) {
root = i;
break;
}
}
//Initialise the count of ideal pairs as zero
idealPairsCount = 0;
// Vector to store the binary indexing of nodes of the tree
vector<long long> bit(N, 0);
// Call the function to find the number of pair of ideal nodes in a given tree
findIdealpairs(root, adj, K, bit, N);
cout<<idealPairsCount<<endl;
return 0;
}

Output:
The number of pair of ideal nodes in the given tree is: 5

Algorithm Complexity:

Time Complexity: O(N * log(N))

In the above program to find the number of pair of ideal nodes in a given tree, findIdealPairs() is recursively being called for each node that takes O(N) time and updating the indexes and querying to get the count of nodes having values in a given range using the Binary indexed tree takes O(log N) time. So, the overall time complexity is O(N * log(N)), where 'N' is the number of nodes in the given tree.

Space Complexity: O(N)

In the above program to find the number of pair of ideal nodes in a given tree, we have created a vector 'bit' to store the indexing of the Binary indexed tree, so the space complexity is O(N), where 'N' is the number of nodes in the given tree.

What is a Binary Indexed Tree? Binary indexed tree also known as fenwick tree is a data structure that stores the elements in an array and helps to calculate the prefix sum efficiently. This data structure is based on the fact that every positive integer can be represented as sum of powers of 2.

Why we have used fenwick tree data structure to solve the above problem? In the above problem, we needed a data structure that can effectively mark an element as an ancestor, remove an element from being an ancestor, and calculate the number of nodes having values in a given range. A fenwick tree is a data structure that can insert an element, remove an element and calculate the sum of elements in a given range efficiently, so we have used a fenwick tree and used its properties to solve the above problem.

Key takeaways-

This article discussed the problem "Find the number of pair of ideal nodes in a given tree", an approach to solve this problem using binary indexed tree, its C++ implementation, and its time and space complexity. If you want to solve more problems on data structure and algorithms for practice, you can visit Coding Ninjas Studio.

Also read, abs in C++ If you think that this blog helped you, then share it with your friends!. Refer to the DSA C++ course for more information.

Until then, All the best for your future endeavors, and Keep Coding.