Find the number of pair of Ideal nodes in a given tree

Given a tree of N nodes and an integer K, each node is numbered between 1 and N. The task is to find the number of pair of ideal nodes in a tree.
A pair of nodes (a, b) is called ideal if

  1. a is an ancestor of b.
  2. And, abs(a – b) ≤ K

Examples:

Input:

K = 2
Output: 4
(1, 2), (1, 3), (3, 4), (3, 5) are such pairs.

Input: Consider the graph in example 1 and k = 3
Output: 6
(1, 2), (1, 3), (1, 4), (3, 4), (3, 5), (3, 6) are such pairs.

Approach: First, we need to traverse the tree using DFS so we need to find the root node, the node without a parent. As we traverse each node we will store it in a data structure to keep track of all the ancestors for the next node. Before doing that, get the number of the node’s ancestors in the range [presentNode – k, presentNode + k] then add it to the total pairs.
We need a data structure which can:

  1. Insert a node as we traverse the tree.
  2. Remove a node as we return.
  3. Give the number of nodes within the range [presentNode – k, PresentNode + k] which were stored.

Binary indexed tree fulfills the above three operations.

We can do the above 3 operations by initializing all the index values of the BIT to 0 and then:

  1. Insert a node by updating the index of that node to 1.
  2. Remove a node by updating the index of that node to 0.
  3. Get the number of similar pairs of the ancestor of that node by querying for the sum of the range [presentNode – k, PresentNode + k]

Below is the implementation of the above approach:

filter_none

edit
close

play_arrow

link
brightness_4
code

// C++ implementation of the approach
#include <bits/stdc++.h>
using namespace std;
#define N 100005
  
int n, k;
  
// Adjacency list
vector<int> al[N];
long long Ideal_pair;
long long bit[N];
bool root_node[N];
  
// bit : bit array
// i and j are starting and 
// ending index INCLUSIVE
long long bit_q(int i, int j)
{
    long long sum = 0ll;
    while (j > 0) {
        sum += bit[j];
        j -= (j & (j * -1));
    }
    i--;
    while (i > 0) {
        sum -= bit[i];
        i -= (i & (i * -1));
    }
    return sum;
}
  
// bit : bit array
// n : size of bit array
// i is the index to be updated
// diff is (new_val - old_val)
void bit_up(int i, long long diff)
{
    while (i <= n) {
        bit[i] += diff;
        i += i & -i;
    }
}
  
// DFS function to find ideal pairs
void dfs(int node)
{
    Ideal_pair += bit_q(max(1, node - k),
                        min(n, node + k));
    bit_up(node, 1);
    for (int i = 0; i < al[node].size(); i++)
        dfs(al[node][i]);
    bit_up(node, -1);
}
  
// Function for initialisation
void initalise()
{
    Ideal_pair = 0;
    for (int i = 0; i <= n; i++) {
        root_node[i] = true;
        bit[i] = 0LL;
    }
}
  
// Function to add an edge
void Add_Edge(int x, int y)
{
    al[x].push_back(y);
    root_node[y] = false;
}
  
// Function to find number of ideal pairs
long long Idealpairs()
{
    // Find root of the tree
    int r = -1;
    for (int i = 1; i <= n; i++)
        if (root_node[i]) {
            r = i;
            break;
        }
  
    dfs(r);
  
    return Ideal_pair;
}
  
// Driver code
int main()
{
    n = 6, k = 3;
  
    initalise();
  
    // Add edges
    Add_Edge(1, 2);
    Add_Edge(1, 3);
    Add_Edge(3, 4);
    Add_Edge(3, 5);
    Add_Edge(3, 6);
  
    // Function call
    cout << Idealpairs();
  
    return 0;
}

chevron_right


Output:

6


My Personal Notes arrow_drop_up

pawanasipugmailcom

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.