Tuesday, November 19, 2024
Google search engine
HomeData Modelling & AIFind distance between two nodes in the given Binary tree for Q...

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

  • 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}

  • As we will compute the lca[][] array we will also calculate 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}

  • 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++




// 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;
}


Java




// Java program to find distance
// between two nodes using LCA
import java.io.*;
import java.util.*;
 
class GFG{
 
static final int MAX = 1000;
// log2(MAX)
static final int log = 10;
 
// Array to store the level
// of each node
static int[] level = new int[MAX];
 
static int[][] lca = new int[MAX][log];
static int[][] dist = new int[MAX][log];
 
// Vector to store tree
@SuppressWarnings("unchecked")
static List<List<int[]> > graph = new ArrayList();
 
static void addEdge(int u, int v, int cost)
{
    graph.get(u).add(new int[]{ v, cost });
    graph.get(v).add(new int[]{ u, cost });
}
 
// Pre-Processing to calculate
// values of lca[][], dist[][]
static 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(int[] i : graph.get(node))
    {
        if (i[0] == parent)
            continue;
             
        dfs(i[0], node, h + 1, i[1]);
    }
}
 
// Function to find the distance
// between given nodes u and v
static 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])
    {
        int temp = u;
        u = v;
        v = temp;
    }
 
    // 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)
    {
        System.out.println(ans);
    }
 
    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];
 
        System.out.println(ans);
    }
}
 
// Driver Code
public static void main(String[] args)
{
     
    // Number of nodes
    int n = 5;
 
    for(int i = 0; i < MAX; i++)
    {
        graph.add(new ArrayList<int[]>());
    }
 
    // 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);
}
}
 
// This code is contributed by jithin


Python3




# Python3 Program to find
# distance between two nodes
# using LCA
MAX = 1000
 
# lg2(MAX)
lg = 10
 
# Array to store the level
# of each node
level = [0 for i in range(MAX)]
 
lca = [[0 for i in range(lg)]
          for j in range(MAX)]
dist = [[0 for i in range(lg)]
           for j in range(MAX)]
 
# Vector to store tree
graph = [[] for i in range(MAX)]
 
def addEdge(u, v, cost):
   
    global graph
     
    graph[u].append([v, cost])
    graph[v].append([u, cost])
 
# Pre-Processing to calculate
# values of lca[][], dist[][]
def dfs(node, parent, h, 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 i in range(1, lg):
        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 i in graph[node]:
        if (i[0] == parent):
            continue
        dfs(i[0], node, h + 1, i[1])
 
# Function to find the distance
# between given nodes u and v
def findDistance(u, v):
   
    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]):
        temp = u
        u = v
        v = temp
 
    # Finding the ancestor of v
    # which is at same level as u
    i = lg - 1
     
    while(i >= 0):
        if (lca[v][i] != -1 and
            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]
             
        i -= 1
 
    # If u is the ancestor of v
    # then u is the LCA of u and v
    if (v == u):
        print(ans)
 
    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
        i = lg - 1
         
        while(i >= 0):
            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]
            i -= 1
 
        # Adding the distance of u and v
        # to its first ancestor
        ans += (dist[u][0] +
                dist[v][0])
 
        print(ans)
 
# Driver Code
if __name__ == '__main__':
   
    # Number of nodes
    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 i in range(1, n + 1):
        for j in range(lg):
            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)
     
# This code is contributed by SURENDRA_GANGWAR


C#




// C# program to find distance
// between two nodes using LCA
using System;
using System.Collections.Generic;
class GFG
{
 
  static readonly int MAX = 1000;
 
  // log2(MAX)
  static readonly int log = 10;
 
  // Array to store the level
  // of each node
  static int[] level = new int[MAX];
  static int[,] lca = new int[MAX,log];
  static int[,] dist = new int[MAX,log];
 
  // List to store tree
  static List<List<int[]> > graph = new List<List<int[]>>();
 
  static void addEdge(int u, int v, int cost)
  {
    graph[u].Add(new int[]{ v, cost });
    graph[v].Add(new int[]{ u, cost });
  }
 
  // Pre-Processing to calculate
  // values of lca[,], dist[,]
  static 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];
      }
    }
 
    foreach(int[] i in graph[node])
    {
      if (i[0] == parent)
        continue;           
      dfs(i[0], node, h + 1, i[1]);
    }
  }
 
  // Function to find the distance
  // between given nodes u and v
  static 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])
    {
      int temp = u;
      u = v;
      v = temp;
    }
 
    // 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)
    {
      Console.WriteLine(ans);
    }
 
    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];
      Console.WriteLine(ans);
    }
  }
 
  // Driver Code
  public static void Main(String[] args)
  {
 
    // Number of nodes
    int n = 5;
    for(int i = 0; i < MAX; i++)
    {
      graph.Add(new List<int[]>());
    }
 
    // 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);
  }
}
 
// This code is contributed by aashish1995


Javascript




<script>
    // Javascript program to find distance
    // between two nodes using LCA
     
    let MAX = 1000;
    // log2(MAX)
    let log = 10;
 
    // Array to store the level
    // of each node
    let level = new Array(MAX);
 
    let lca = new Array(MAX);
    let dist = new Array(MAX);
 
    // Vector to store tree
    let graph = [];
 
    function addEdge(u, v, cost)
    {
        graph[u].push([ v, cost ]);
        graph[v].push([ u, cost ]);
    }
 
    // Pre-Processing to calculate
    // values of lca[][], dist[][]
    function dfs(node, parent, h, 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(let 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(let i = 0; i < graph[node].length; i++)
        {
            if (graph[node][i][0] == parent)
                continue;
 
            dfs(graph[node][i][0], node, h + 1, graph[node][i][1]);
        }
    }
 
    // Function to find the distance
    // between given nodes u and v
    function findDistance(u, v)
    {
        let 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])
        {
            let temp = u;
            u = v;
            v = temp;
        }
 
        // Finding the ancestor of v
        // which is at same level as u
        for(let 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)
        {
            document.write(ans + "</br>");
        }
 
        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(let 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];
 
            document.write(ans + "</br>");
        }
    }
     
    // Number of nodes
    let n = 5;
  
    for(let i = 0; i < MAX; i++)
    {
        graph.push([]);
    }
  
    // 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(let i = 1; i <= n; i++)
    {
        lca[i] = new Array(log);
        dist[i] = new Array(log);
        for(let 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);
 
// This code is contributed by decode2207.
</script>


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

Space Complexity: O(N*log(N)) 
We are storing the LCA and distance of all the nodes in two 2-D arrays.
 

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!

RELATED ARTICLES

Most Popular

Recent Comments