CSES Solutions – Tree Distances I
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:
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++ 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;
}
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
# 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()
// 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.
Contact Us