Find the Number of Pair of Ideal Nodes in a Given 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:

1. Node - The current node being considered (start from the root node)
2. adj - The adjacency list representation of the given tree
3. K - Given value of K
4. bit - vector to store the indexing of the tree
5. 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.

Check out this problem - Pair Sum In Array.

FAQs

1. 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.
    
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.