Categories
28°C, New York
Home Data Structures and Algorithms CSES Solutions – Tree Distances II

CSES Solutions – Tree Distances II

Find distance of nodes from root in a tree for multiple queries

Find Pattern in Infinite Binary Stream I

You are given a tree consisting of n nodes. Your task is to determine for each node the sum of the distances from the node to all other nodes.

Example:

Input: edges = { { 1, 2 }, { 1, 3 }, { 3, 4 }, { 3, 5 } }
Output: 6 9 5 8 8

Input: edges = { { 1, 2 }, { 2, 3 }, { 3, 4 }, { 3, 5 } }
Output: 9 6 5 8 8

Approach:

Lets see the intuition into steps

  1. Finding the sum of distances from a single node: If we root the tree at a particular node, we can use Depth-First Search (DFS) to find the depth of each other node. The sum of these depths gives us the sum of distances from the root node to all other nodes.
  2. Finding the sum of distances from all nodes: Doing the above for each node would be too slow if n is large. So, we need a faster way. The key is to use the answer for one node to quickly find the answer for its neighbors.
  3. Transitioning from one node to its neighbor: If we reroot the tree at a neighbor of the current root, the depths of all nodes in the new root’s subtree decrease by 1, and the depths of all nodes outside of its subtree increase by 1. This means the change in the sum of distances is exactly n – 2 * (size of new root’s subtree).
  4. Calculating the answer for all nodes: We can use DFS twice to find the answer for all nodes. The first DFS finds the answer for one node and the size of each node’s subtree. The second DFS computes the answer for all nodes using the formula from step 3.

Steps-by-step approach:

  • DFS1 Function:
    • Performs a depth-first search (DFS) starting from node 1.
    • Calculates the answer for node 1 and initializes the size of each node’s subtree.
    • For each child of the current node:
      • If the child is not the parent, it recursively calls dfs1() on the child, passing the child as the current node, the current node as the parent, and the depth incremented by 1.
      • It updates dp[node] by adding the size of the child’s subtree.
  • DFS2 Function:
    • Performs a depth-first search (DFS) starting from node 1.
    • For each child of the current node:
      • If the child is not the parent, it calculates the answer for the child using the formula provided in the comments.
      • It recursively calls dfs2() on the child, passing the child as the current node and the current node as the parent.
  • Main Function:
    • Initializes the number of nodes n and the edges of the tree.
    • Builds the graph by adding edges to the adjacency list representation.
    • Calls dfs1() to compute initial values.
    • Calls dfs2() to calculate the final answer for each node.
    • Prints the answer for each node.

Below is the implementation of the above approach:

C++ 
#include <bits/stdc++.h>
typedef long long ll;
using namespace std;

int n;
vector<int> graph[200001];
ll dp[200001], ans[200001];

// First DFS to calculate the answer for node 1 and the size
// of each node's subtree
void dfs1(int node = 1, int parent = 0, ll depth = 0)
{
    // Add depth to the answer for node 1
    ans[1] += depth;

    // Initialize the size of the subtree rooted at node to
    // 1
    dp[node] = 1;

    // For each child of node
    for (int i : graph[node])

        // If the child is not the parent
        if (i != parent) {

            // Recurse on the child
            dfs1(i, node, depth + 1);

            // Add the size of the child's subtree to
            // dp[node]
            dp[node] += dp[i];
        }
}

// Second DFS to calculate the answer for all nodes
void dfs2(int node = 1, int parent = 0)
{

    // For each child of node
    for (int i : graph[node])

        // If the child is not the parent
        if (i != parent) {

            // Calculate the answer for the child
            ans[i] = ans[node] + n - 2 * dp[i];

            // Recurse on the child
            dfs2(i, node);
        }
}

// Driver code
int main()
{

    n = 5;
    vector<pair<int, int> > edges
        = { { 1, 2 }, { 1, 3 }, { 3, 4 }, { 3, 5 } };

    // Build the graph
    for (auto edge : edges) {
        int a = edge.first, b = edge.second;
        graph[a].push_back(b);
        graph[b].push_back(a);
    }

    // Run the first DFS
    dfs1();

    // Run the second DFS
    dfs2();

    // Print the answer for each node
    for (int i = 1; i <= n; i++)
        cout << ans[i] << ' ';

    return 0;
}
Java 
import java.util.ArrayList;
import java.util.List;

public class Main {
    static int n;
    static List<Integer>[] graph;
    static long[] dp, ans;

    // First DFS to calculate the answer for node 1 and the
    // size of each node's subtree
    static void dfs1(int node, int parent, long depth)
    {
        // Add depth to the answer for node 1
        ans[1] += depth;

        // Initialize the size of the subtree rooted at node
        // to 1
        dp[node] = 1;

        // For each child of node
        for (int i : graph[node]) {
            // If the child is not the parent
            if (i != parent) {
                // Recurse on the child
                dfs1(i, node, depth + 1);
                // Add the size of the child's subtree to
                // dp[node]
                dp[node] += dp[i];
            }
        }
    }

    // Second DFS to calculate the answer for all nodes
    static void dfs2(int node, int parent)
    {
        // For each child of node
        for (int i : graph[node]) {
            // If the child is not the parent
            if (i != parent) {
                // Calculate the answer for the child
                ans[i] = ans[node] + n - 2 * dp[i];
                // Recurse on the child
                dfs2(i, node);
            }
        }
    }

    // Driver code
    public static void main(String[] args)
    {
        n = 5;
        graph = new ArrayList[n + 1];
        for (int i = 1; i <= n; i++) {
            graph[i] = new ArrayList<>();
        }

        List<int[]> edges = List.of(
            new int[] { 1, 2 }, new int[] { 1, 3 },
            new int[] { 3, 4 }, new int[] { 3, 5 });

        // Build the graph
        for (int[] edge : edges) {
            int a = edge[0], b = edge[1];
            graph[a].add(b);
            graph[b].add(a);
        }

        dp = new long[n + 1];
        ans = new long[n + 1];

        // Run the first DFS
        dfs1(1, 0, 0);

        // Run the second DFS
        dfs2(1, 0);

        // Print the answer for each node
        for (int i = 1; i <= n; i++)
            System.out.print(ans[i] + " ");
    }
}

// This code is contributed by shivamgupta310570
Python 
from collections import defaultdict

# Initialize variables
n = 5
graph = defaultdict(list)
dp = [0] * (n+1)
ans = [0] * (n+1)

# First DFS to calculate the answer for node 1 and the size of each node's subtree


def dfs1(node=1, parent=0, depth=0):
    # Add depth to the answer for node 1
    ans[1] += depth

    # Initialize the size of the subtree rooted at node to 1
    dp[node] = 1

    # For each child of node
    for i in graph[node]:
        # If the child is not the parent
        if i != parent:
            # Recurse on the child
            dfs1(i, node, depth + 1)

            # Add the size of the child's subtree to dp[node]
            dp[node] += dp[i]

# Second DFS to calculate the answer for all nodes


def dfs2(node=1, parent=0):
    # For each child of node
    for i in graph[node]:
        # If the child is not the parent
        if i != parent:
            # Calculate the answer for the child
            ans[i] = ans[node] + n - 2 * dp[i]

            # Recurse on the child
            dfs2(i, node)


# Driver code
if __name__ == "__main__":
    # Define edges
    edges = [(1, 2), (1, 3), (3, 4), (3, 5)]

    # Build the graph
    for edge in edges:
        a, b = edge
        graph[a].append(b)
        graph[b].append(a)

    # Run the first DFS
    dfs1()

    # Run the second DFS
    dfs2()

    # Print the answer for each node
    for i in range(1, n+1):
        print(ans[i], end=' ')
JavaScript 
let n;
let graph = [];
let dp = [];
let ans = [];

// First DFS to calculate the answer for node 1 and the size
// of each node's subtree
function dfs1(node = 1, parent = 0, depth = 0) {
    // Add depth to the answer for node 1
    ans[1] += depth;

    // Initialize the size of the subtree rooted at node to 1
    dp[node] = 1;

    // For each child of node
    for (let i of graph[node]) {
        // If the child is not the parent
        if (i !== parent) {
            // Recurse on the child
            dfs1(i, node, depth + 1);

            // Add the size of the child's subtree to dp[node]
            dp[node] += dp[i];
        }
    }
}

// Second DFS to calculate the answer for all nodes
function dfs2(node = 1, parent = 0) {
    // For each child of node
    for (let i of graph[node]) {
        // If the child is not the parent
        if (i !== parent) {
            // Calculate the answer for the child
            ans[i] = ans[node] + n - 2 * dp[i];

            // Recurse on the child
            dfs2(i, node);
        }
    }
}

// Driver code
function main() {
    n = 5;
    let edges = [[1, 2], [1, 3], [3, 4], [3, 5]];

    // Initialize graph
    for (let i = 0; i <= n; i++) {
        graph.push([]);
    }

    // Build the graph
    for (let edge of edges) {
        let a = edge[0], b = edge[1];
        graph[a].push(b);
        graph[b].push(a);
    }

    // Initialize dp and ans arrays
    for (let i = 0; i <= n; i++) {
        dp.push(0);
        ans.push(0);
    }

    // Run the first DFS
    dfs1();

    // Run the second DFS
    dfs2();

    // Print the answer for each node
    for (let i = 1; i <= n; i++) {
        process.stdout.write(ans[i] + ' ');
    }
}

main();

Output
6 9 5 8 8

Time Complexity: O(n)
Auxiliary Space: O(n)



Tags:
#CSES Problems #DSA #Tree #Tree

Find distance of nodes from root in a tree for multiple queries

Find Pattern in Infinite Binary Stream I