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:

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.