Open In App

CSES Solutions – Tree Distances I

Last Updated : 10 Apr, 2024
Suggest changes
Share
Like Article
Like
Report

Given a tree consisting of n nodes (1,2,...,n) and n-1 lines describing the edges between nodes u and v. The task is to determine for each node the maximum distance to another node.

Example:

Input: n= 5, v= {{1,2}, {1,3}, {3,4}, {3,5}}

Output: 2 3 2 3 3

Explanation:

Screenshot-2024-03-11-170227
Tree

The farthest node from node 1 is node 4, with a distance of 2.
The farthest node from node 2 is node 4, with a distance of 3.
The farthest node from node 3 is node 2, with a distance of 2.
The farthest node from node 4 is node 2, with a distance of 3.
The farthest node from node 5 is node 2, with a distance of 3.

Approach: To solve the problem, follow the idea below:

The idea is to find the diameter (maximum distance between any two nodes) of the tree and use it to calculate the maximum distance for each node to any other node. Consider any node u in the tree. The farthest node from u must be the either end of the diameter of the tree. This is because any path from u to a node not on the diameter could be made longer by extending it to a node on the diameter. Use BFS to find the farthest node from 1 and store it in di_st. Again use BFS to find the farthest node from di_st and store it in di_end, these nodes are farthest apart. Then use DFS to calculate the distance of each node from di_st and di_end. The maximum distance for a node to any other node is the maximum of these two distances.

Step-by-step algorithm:

  • Create an adjacency list adj to represent the tree.
  • Use BFS to find the farthest node from any node, take 1 and store it in di_st.
  • Again use BFS to find the farthest node from di_st and store it in di_end.
  • Use DFS to calculate the distance of each node from di_st.
  • Store the distance of the node from the starting node in the first column of vector dp.
  • Recursively call for the child node and increase the distance by 1.
  • Again use DFS to calculate the distance of each node from di_end.
  • Store the distance of the node from the starting node in the second column of vector dp.
  • Recursively call for the child node and increase the distance by 1.
  • For each node, find the maximum distance by taking the maximum of the distance calculated from dt_st and distance calculated from dt_end.
  • Return the maximum distances in a vector mxdist.

Below is the implementation of the above algorithm:

C++
// C++ code #include <bits/stdc++.h> using namespace std; // Depth-first search function to calculate the distance of // each node from a given node void dfs(int parent, int node, int dist, int c,  vector<int> adj[], vector<vector<int> >& dp) {  // Store the distance of the node from the starting node  dp[node][c] = dist;  // Iterate over all the adjacent nodes  for (auto child : adj[node])  // If the adjacent node is not the parent  if (child != parent)  // Recursively call for the child node and  // increase the distance by 1  dfs(node, child, dist + 1, c, adj, dp); } // Breadth-first search function to find the node at the // maximum distance from a given node int bfs(int node, int n, vector<int> adj[]) {  // Queue to store the node and its distance from the  // starting node  queue<pair<int, int> > q;  // Push the starting node into the queue with distance 0  q.push({ node, 0 });  // Visited array to keep track of visited nodes  vector<bool> vis(n + 1, false);  // Mark the starting node as visited  vis[node] = true;  int u, dist;  while (!q.empty()) {  // Node at the front of the queue  u = q.front().first;  // Distance of the node from the starting node  dist = q.front().second;  // Remove the node from the queue  q.pop();  // Iterate over all the adjacent nodes  for (auto child : adj[u]) {  if (!vis[child]) {  // Push the node into the queue  q.push({ child, dist + 1 });  // Mark the node as visited  vis[child] = true;  }  }  }  // Return the node at the maximum distance  return u; } vector<int> treedistance1(int n, vector<vector<int> >& v) {  vector<int> adj[n + 1];  for (auto i : v) {  adj[i[0]].push_back(i[1]);  adj[i[1]].push_back(i[0]);  }  vector<int> mxdist;  // Node at the maximum distance from node 1  int di_st = bfs(1, n, adj);  // Node at the maximum distance from node di_st  int di_end = bfs(di_st, n, adj);  // 2D vector to store the maximum distance of each node  vector<vector<int> > dp(n + 1, vector<int>(2, 0));  // DFS from di_st  dfs(0, di_st, 0, 0, adj, dp);  // DFS from di_end  dfs(0, di_end, 0, 1, adj, dp);  // Take the max distance  for (int i = 1; i <= n; i++)  mxdist.push_back(max(dp[i][0], dp[i][1]));  return mxdist; } // Driver Code int main() {  int n = 5;  vector<vector<int> > v  = { { 1, 2 }, { 1, 3 }, { 3, 4 }, { 3, 5 } };  vector<int> mxdist = treedistance1(n, v);  if( n == 1 ){  cout<<0<<endl;  }  for (auto i : mxdist)  cout << i << " ";  return 0; } 
Java
import java.util.*; public class TreeDistance {  // Depth-first search function to calculate the distance of  // each node from a given node  static void dfs(int parent, int node, int dist, int c,  ArrayList<Integer>[] adj, int[][] dp) {  // Store the distance of the node from the starting node  dp[node][c] = dist;  // Iterate over all the adjacent nodes  for (int child : adj[node]) {  // If the adjacent node is not the parent  if (child != parent) {  // Recursively call for the child node and  // increase the distance by 1  dfs(node, child, dist + 1, c, adj, dp);  }  }  }  // Breadth-first search function to find the node at the  // maximum distance from a given node  static int bfs(int node, int n, ArrayList<Integer>[] adj) {  // Queue to store the node and its distance from the  // starting node  Queue<int[]> q = new LinkedList<>();  // Push the starting node into the queue with distance 0  q.offer(new int[]{node, 0});  // Visited array to keep track of visited nodes  boolean[] vis = new boolean[n + 1];  // Mark the starting node as visited  vis[node] = true;  int u = 0, dist = 0;  while (!q.isEmpty()) {  // Node at the front of the queue  int[] curr = q.poll();  u = curr[0];  // Distance of the node from the starting node  dist = curr[1];  // Iterate over all the adjacent nodes  for (int child : adj[u]) {  if (!vis[child]) {  // Push the node into the queue  q.offer(new int[]{child, dist + 1});  // Mark the node as visited  vis[child] = true;  }  }  }  // Return the node at the maximum distance  return u;  }  public static ArrayList<Integer> treeDistance1(int n, ArrayList<int[]> v) {  ArrayList<Integer>[] adj = new ArrayList[n + 1];  for (int i = 1; i <= n; i++) {  adj[i] = new ArrayList<>();  }  for (int[] edge : v) {  adj[edge[0]].add(edge[1]);  adj[edge[1]].add(edge[0]);  }  ArrayList<Integer> mxdist = new ArrayList<>();  // Node at the maximum distance from node 1  int di_st = bfs(1, n, adj);  // Node at the maximum distance from node di_st  int di_end = bfs(di_st, n, adj);  // 2D array to store the maximum distance of each node  int[][] dp = new int[n + 1][2];  // DFS from di_st  dfs(0, di_st, 0, 0, adj, dp);  // DFS from di_end  dfs(0, di_end, 0, 1, adj, dp);  // Take the max distance  for (int i = 1; i <= n; i++)  mxdist.add(Math.max(dp[i][0], dp[i][1]));  return mxdist;  }  // Driver Code  public static void main(String[] args) {  int n = 5;  ArrayList<int[]> v = new ArrayList<>(Arrays.asList(  new int[]{1, 2}, new int[]{1, 3}, new int[]{3, 4}, new int[]{3, 5}  ));  ArrayList<Integer> mxdist = treeDistance1(n, v);  if (n == 1) {  System.out.println(0);  }  for (int i : mxdist)  System.out.print(i + " ");  } } // This code is contributed by shivamgupta0987654321 
Python3
# Importing the required libraries from collections import deque # Depth-first search function to calculate the distance of # each node from a given node def dfs(parent, node, dist, c, adj, dp): # Store the distance of the node from the starting node dp[node][c] = dist # Iterate over all the adjacent nodes for child in adj[node]: # If the adjacent node is not the parent if child != parent: # Recursively call for the child node and # increase the distance by 1 dfs(node, child, dist + 1, c, adj, dp) # Breadth-first search function to find the node at the # maximum distance from a given node def bfs(node, n, adj): # Queue to store the node and its distance from the # starting node q = deque([(node, 0)]) # Visited array to keep track of visited nodes vis = [False] * (n + 1) # Mark the starting node as visited vis[node] = True while q: # Node at the front of the queue u, dist = q.popleft() # Iterate over all the adjacent nodes for child in adj[u]: if not vis[child]: # Push the node into the queue q.append((child, dist + 1)) # Mark the node as visited vis[child] = True # Return the node at the maximum distance return u def treedistance1(n, v): adj = [[] for _ in range(n + 1)] for i in v: adj[i[0]].append(i[1]) adj[i[1]].append(i[0]) mxdist = [] # Node at the maximum distance from node 1 di_st = bfs(1, n, adj) # Node at the maximum distance from node di_st di_end = bfs(di_st, n, adj) # 2D vector to store the maximum distance of each node dp = [[0, 0] for _ in range(n + 1)] # DFS from di_st dfs(0, di_st, 0, 0, adj, dp) # DFS from di_end dfs(0, di_end, 0, 1, adj, dp) # Take the max distance for i in range(1, n + 1): mxdist.append(max(dp[i][0], dp[i][1])) return mxdist # Driver Code def main(): n = 5 v = [[1, 2], [1, 3], [3, 4], [3, 5]] mxdist = treedistance1(n, v) if n == 1: print(0) for i in mxdist: print(i, end=" ") if __name__ == "__main__": main() 
JavaScript
// Depth-first search function to calculate the distance of // each node from a given node function dfs(parent, node, dist, c, adj, dp) {  // Store the distance of the node from the starting node  dp[node][c] = dist;  // Iterate over all the adjacent nodes  for (let child of adj[node]) {  // If the adjacent node is not the parent  if (child !== parent) {  // Recursively call for the child node and  // increase the distance by 1  dfs(node, child, dist + 1, c, adj, dp);  }  } } // Breadth-first search function to find the node at the // maximum distance from a given node function bfs(node, n, adj) {  // Queue to store the node and its distance from the  // starting node  let q = [];  // Push the starting node into the queue with distance 0  q.push([node, 0]);  // Visited array to keep track of visited nodes  let vis = new Array(n + 1).fill(false);  // Mark the starting node as visited  vis[node] = true;  let u, dist;  while (q.length !== 0) {  // Node at the front of the queue  [u, dist] = q.shift();  // Iterate over all the adjacent nodes  for (let child of adj[u]) {  if (!vis[child]) {  // Push the node into the queue  q.push([child, dist + 1]);  // Mark the node as visited  vis[child] = true;  }  }  }  // Return the node at the maximum distance  return u; } function treedistance1(n, v) {  let adj = new Array(n + 1).fill(null).map(() => []);  for (let i of v) {  adj[i[0]].push(i[1]);  adj[i[1]].push(i[0]);  }  let mxdist = [];  // Node at the maximum distance from node 1  let di_st = bfs(1, n, adj);  // Node at the maximum distance from node di_st  let di_end = bfs(di_st, n, adj);  // 2D array to store the maximum distance of each node  let dp = new Array(n + 1).fill(null).map(() => [0, 0]);  // DFS from di_st  dfs(0, di_st, 0, 0, adj, dp);  // DFS from di_end  dfs(0, di_end, 0, 1, adj, dp);  // Take the max distance  for (let i = 1; i <= n; i++)  mxdist.push(Math.max(dp[i][0], dp[i][1]));  return mxdist; } // Driver Code function main() {  let n = 5;  let v = [[1, 2], [1, 3], [3, 4], [3, 5]];  let mxdist = treedistance1(n, v);  let output = ""; // Initialize an empty string to store the output  if (n === 1) {  output += "0 "; // Append "0" to the output string  } else {  for (let i of mxdist)  output += i + " "; // Append each distance followed by a space to the output string  }  console.log(output.trim()); // Print the output string after trimming any leading/trailing spaces } // Call the main function main(); 

Output
2 3 2 3 3 

Time Complexity: O(n), where n is the number of node.
Auxiliary Space: O(n), where n is the number of node.



Similar Reads