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 pairs of ideal nodes in a tree.
A pair of nodes (a, b) is called ideal if
- a is an ancestor of b.
- 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:
- Insert a node as we traverse the tree.
- Remove a node as we return.
- 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:
- Insert a node by updating the index of that node to 1.
- Remove a node by updating the index of that node to 0.
- 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:
C++
// C++ implementation of the approach#include <bits/stdc++.h>using namespace std;#define N 100005int n, k;// Adjacency listvector<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 INCLUSIVElong 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 pairsvoid 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 initialisationvoid initialise(){ Ideal_pair = 0; for (int i = 0; i <= n; i++) { root_node[i] = true; bit[i] = 0LL; }}// Function to add an edgevoid Add_Edge(int x, int y){ al[x].push_back(y); root_node[y] = false;}// Function to find number of ideal pairslong 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 codeint main(){ n = 6, k = 3; initialise(); // 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;} |
Java
// Java implementation of the approach import java.util.*;class GFG{ static final int N = 100005;static int n, k; // Adjacency list @SuppressWarnings("unchecked")static Vector<Integer> []al = new Vector[N]; static long Ideal_pair; static long []bit = new long[N]; static boolean []root_node = new boolean[N]; // bit : bit array // i and j are starting and // ending index INCLUSIVE static long bit_q(int i, int j) { long sum = 0; 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) static void bit_up(int i, long diff) { while (i <= n) { bit[i] += diff; i += i & -i; } } // DFS function to find ideal pairs static void dfs(int node) { Ideal_pair += bit_q(Math.max(1, node - k), Math.min(n, node + k)); bit_up(node, 1); for(int i = 0; i < al[node].size(); i++) dfs(al[node].get(i)); bit_up(node, -1); } // Function for initialisation static void initialise() { Ideal_pair = 0; for (int i = 0; i <= n; i++) { root_node[i] = true; bit[i] = 0; } } // Function to add an edge static void Add_Edge(int x, int y) { al[x].add(y); root_node[y] = false; } // Function to find number of ideal pairs static 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 public static void main(String[] args) { n = 6; k = 3; for(int i = 0; i < al.length; i++) al[i] = new Vector<Integer>(); initialise(); // Add edges Add_Edge(1, 2); Add_Edge(1, 3); Add_Edge(3, 4); Add_Edge(3, 5); Add_Edge(3, 6); // Function call System.out.print(Idealpairs()); }} // This code is contributed by Amit Katiyar |
Python3
# Python3 implementation of the approachN = 100005Ideal_pair = 0# Adjacency listal = [[] for i in range(100005)]bit = [0 for i in range(N)]root_node = [0 for i in range(N)]# bit : bit array# i and j are starting and # ending index INCLUSIVEdef bit_q(i, j): sum = 0 while (j > 0): sum += bit[j] j -= (j & (j * -1)) i -= 1 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)def bit_up(i, diff): while (i <= n): bit[i] += diff i += i & -i# DFS function to find ideal pairsdef dfs(node, x): Ideal_pair = x Ideal_pair += bit_q(max(1, node - k), min(n, node + k)) bit_up(node, 1) for i in range(len(al[node])): Ideal_pair = dfs(al[node][i], Ideal_pair) bit_up(node, -1) return Ideal_pair # Function for initialisationdef initialise(): Ideal_pair = 0; for i in range(n + 1): root_node[i] = True bit[i] = 0# Function to add an edgedef Add_Edge(x, y): al[x].append(y) root_node[y] = False# Function to find number of ideal pairsdef Idealpairs(): # Find root of the tree r = -1 for i in range(1, n + 1, 1): if (root_node[i]): r = i break Ideal_pair = dfs(r, 0) return Ideal_pair# Driver codeif __name__ == '__main__': n = 6 k = 3 initialise() # Add edges Add_Edge(1, 2) Add_Edge(1, 3) Add_Edge(3, 4) Add_Edge(3, 5) Add_Edge(3, 6) # Function call print(Idealpairs())# This code is contributed by# Surendra_Gangwar |
C#
// C# implementation of the // above approach using System;using System.Collections.Generic;class GFG{ static readonly int N = 100005;static int n, k; // Adjacency list static List<int> []al = new List<int>[N]; static long Ideal_pair; static long []bit = new long[N]; static bool []root_node = new bool[N]; // bit : bit array // i and j are starting and // ending index INCLUSIVE static long bit_q(int i, int j) { long sum = 0; 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) static void bit_up(int i, long diff) { while (i <= n) { bit[i] += diff; i += i & -i; } } // DFS function to find // ideal pairs static void dfs(int node) { Ideal_pair += bit_q(Math.Max(1, node - k), Math.Min(n, node + k)); bit_up(node, 1); for(int i = 0; i < al[node].Count; i++) dfs(al[node][i]); bit_up(node, -1); } // Function for // initialisation static void initialise() { Ideal_pair = 0; for (int i = 0; i <= n; i++) { root_node[i] = true; bit[i] = 0; } } // Function to add an edge static void Add_Edge(int x, int y) { al[x].Add(y); root_node[y] = false; } // Function to find number // of ideal pairs static 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 public static void Main(String[] args) { n = 6; k = 3; for(int i = 0; i < al.Length; i++) al[i] = new List<int>(); initialise(); // Add edges Add_Edge(1, 2); Add_Edge(1, 3); Add_Edge(3, 4); Add_Edge(3, 5); Add_Edge(3, 6); // Function call Console.Write(Idealpairs()); }} // This code is contributed by Amit Katiyar |
Javascript
<script>// Javascript implementation of the approachlet N = 100005;let n, k;// Adjacency listlet al = new Array(N).fill(0).map((t) => []);let Ideal_pair;let bit = new Array(N);let root_node = new Array(N);// bit : bit array// i and j are starting and// ending index INCLUSIVEfunction bit_q(i, j) { let sum = 0; 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)function bit_up(i, diff) { while (i <= n) { bit[i] += diff; i += i & -i; }}// DFS function to find ideal pairsfunction dfs(node) { Ideal_pair += bit_q(Math.max(1, node - k), Math.min(n, node + k)); bit_up(node, 1); for (let i = 0; i < al[node].length; i++) dfs(al[node][i]); bit_up(node, -1);}// Function for initialisationfunction initialise() { Ideal_pair = 0; for (let i = 0; i <= n; i++) { root_node[i] = true; bit[i] = 0; }}// Function to add an edgefunction Add_Edge(x, y) { al[x].push(y); root_node[y] = false;}// Function to find number of ideal pairsfunction Idealpairs() { // Find root of the tree let r = -1; for (let i = 1; i <= n; i++) if (root_node[i]) { r = i; break; } dfs(r); return Ideal_pair;}// Driver coden = 6, k = 3;initialise();// Add edgesAdd_Edge(1, 2);Add_Edge(1, 3);Add_Edge(3, 4);Add_Edge(3, 5);Add_Edge(3, 6);// Function calldocument.write(Idealpairs());// This code is contributed by _saurabh_jaiswal</script> |
Output:
6
Time Complexity: O(N*log(N))
Auxiliary Space: O(N)
Feeling lost in the world of random DSA topics, wasting time without progress? It’s time for a change! Join our DSA course, where we’ll guide you on an exciting journey to master DSA efficiently and on schedule.
Ready to dive in? Explore our Free Demo Content and join our DSA course, trusted by over 100,000 neveropen!
Ready to dive in? Explore our Free Demo Content and join our DSA course, trusted by over 100,000 neveropen!
