CSES Solutions – Tree Distances II
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 8Input: edges = { { 1, 2 }, { 2, 3 }, { 3, 4 }, { 3, 5 } }
Output: 9 6 5 8 8
Approach:
Lets see the intuition into steps
- 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.
- 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.
- 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).
- 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:
#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;
}
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
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=' ')
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)
Contact Us