Find distance between two nodes in the given Binary tree for Q queries

Given a binary tree having N nodes and weight of N-1 edges. The distance between two nodes is the sum of the weight of edges on the path between two nodes. Each query contains two integers U and V, the task is to find the distance between node U and V.

Examples:

Input:

Output: 3 5 12 12
Explanation:
Distance between nodes 1 to 3 = weight(1, 3) = 2
Distance between nodes 2 to 3 = weight(1, 2) + weight(1, 3) = 5
Distance between nodes 3 to 5 = weight(1, 3) + weight(1, 2) + weight(2, 5) = 12
Distance between nodes 4 to 5 = weight(4, 2) + weight(2, 5) = 12



Approach: The idea is to use LCA in a tree using Binary Lifting Technique.

  1. Binary Lifting is a Dynamic Programming approach where we pre-compute an array lca[i][j] where i = [1, n], j = [1, log(n)] and lca[i][j] contains 2j-th ancestor of node i.
    • For computing the values of lca[][], the following recursion may be used
      lca[i][j] =\begin{cases} parent[i] & \text{ ;if } j=0 \\ lca[lca[i][j - 1]][j - 1]  & \text{ ;if } j>0 \end{cases}
  2. As we will compute the lca[][] array we will also calulate the distance[][] where distance[i][j] contains the distance from node i to its 2j-th ancestor
    • For computing the values of dist[][], the following recursion may be used.
      dist[i][j] =\begin{cases} cost(i, parent[i]) & \text{ ;if } j=0  \\ dist[i][j] = dist[i][j - 1] + dist[lca[i][j - 1]][j - 1];   & \text{ ;if } j>0 \end{cases}
  3. After precomputation we find the distance between (u, v) as we find the least common ancestor of (u, v).

Below is the implementation of the above approach:

C++

filter_none

edit
close

play_arrow

link
brightness_4
code

// C++ Program to find distance
// between two nodes using LCA
  
#include <bits/stdc++.h>
using namespace std;
  
#define MAX 1000
  
#define log 10 // log2(MAX)
  
// Array to store the level
// of each node
int level[MAX];
  
int lca[MAX][log];
int dist[MAX][log];
  
// Vector to store tree
vector<pair<int, int> > graph[MAX];
  
void addEdge(int u, int v, int cost)
{
    graph[u].push_back({ v, cost });
    graph[v].push_back({ u, cost });
}
  
// Pre-Processing to calculate
// values of lca[][], dist[][]
void dfs(int node, int parent,
         int h, int cost)
{
    // Using recursion formula to
    // calculate the values
    // of lca[][]
    lca[node][0] = parent;
  
    // Storing the level of
    // each node
    level[node] = h;
    if (parent != -1) {
        dist[node][0] = cost;
    }
  
    for (int i = 1; i < log; i++) {
        if (lca[node][i - 1] != -1) {
  
            // Using recursion formula to
            // calculate the values of
            // lca[][] and dist[][]
            lca[node][i]
                = lca[lca[node]
                         [i - 1]]
                     [i - 1];
  
            dist[node][i]
                = dist[node][i - 1]
                  + dist[lca[node][i - 1]]
                        [i - 1];
        }
    }
  
    for (auto i : graph[node]) {
        if (i.first == parent)
            continue;
        dfs(i.first, node, 
h + 1, i.second);
    }
}
  
// Function to find the distance
// between given nodes u and v
void findDistance(int u, int v)
{
  
    int ans = 0;
  
    // The node which is present
    // farthest from the root node
    // is taken as v. If u is
    // farther from root node
    // then swap the two
    if (level[u] > level[v])
        swap(u, v);
  
    // Finding the ancestor of v
    // which is at same level as u
    for (int i = log - 1; i >= 0; i--) {
  
        if (lca[v][i] != -1
            && level[lca[v][i]]
                   >= level[u]) {
  
            // Adding distance of node
            // v till its 2^i-th ancestor
            ans += dist[v][i];
            v = lca[v][i];
        }
    }
  
    // If u is the ancestor of v
    // then u is the LCA of u and v
    if (v == u) {
  
        cout << ans << endl;
    }
  
    else {
  
        // Finding the node closest to the
        // root which is not the common
        // ancestor of u and v i.e. a node
        // x such that x is not the common
        // ancestor of u and v but lca[x][0] is
        for (int i = log - 1; i >= 0; i--) {
  
            if (lca[v][i] != lca[u][i]) {
  
                // Adding the distance
                // of v and u to
                // its 2^i-th ancestor
                ans += dist[u][i] + dist[v][i];
  
                v = lca[v][i];
                u = lca[u][i];
            }
        }
  
        // Adding the distance of u and v
        // to its first ancestor
        ans += dist[u][0] + dist[v][0];
  
        cout << ans << endl;
    }
}
  
// Driver Code
int main()
{
  
    // Number of nodes
    int n = 5;
  
    // Add edges with their cost
    addEdge(1, 2, 2);
    addEdge(1, 3, 3);
    addEdge(2, 4, 5);
    addEdge(2, 5, 7);
  
    // Initialising lca and dist values
    // with -1 and 0 respectively
    for (int i = 1; i <= n; i++) {
        for (int j = 0; j < log; j++) {
            lca[i][j] = -1;
            dist[i][j] = 0;
        }
    }
  
    // Perform DFS
    dfs(1, -1, 0, 0);
  
    // Query 1: {1, 3}
    findDistance(1, 3);
  
    // Query 2: {2, 3}
    findDistance(2, 3);
  
    // Query 3: {3, 5}
    findDistance(3, 5);
  
    return 0;
}

chevron_right


Output:

3
5
12

Time Complexity: The time taken in pre-processing is O(N logN) and every query takes O(logN) time. Therefore, overall time complexity of the solution is O(N logN).

competitive-programming-img




My Personal Notes arrow_drop_up

Check out this Author's contributed articles.

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.