Station Traffic
National Olympiad — Regional round (2022 season)
Submission name: stations.cpp
Problem summary
For every vertex of a weighted tree, sum the lengths of all unordered simple paths passing through that vertex, including paths ending there.
Input
N, then N−1 edges x y length.
Output
The N station scores, one per line.
Constraints
1 ≤ N ≤ 100000; 1 ≤ edge length ≤ 1000.
Worked example
An independently prepared example for this explanation.
Input
6 1 2 3 1 4 5 3 4 7 4 5 9 4 6 10
Output
107 61 67 182 75 79
Solution analysis
- Root the tree. Compute each subtree size and down[v], the sum of distances from v to vertices in its subtree, using a reverse iterative traversal. Then reroot distance sums: all[child]=all[parent]+length*(N−2*size[child]).
- Removing v partitions the other vertices into neighbor branches. For a branch containing s vertices with total distance d from v, its distance contribution appears once in paths ending at v, and once for each endpoint outside that branch. Its total contribution is therefore d*(N−s). Sum this over all branches.
- For sample station 1, the branch through 2 has size 1 and distance sum 3, while the branch through 4 has size 4 and distance sum 46. The score is 3*(6−1)+46*(6−4)=107. For a leaf, its sole branch has N−1 vertices, so its score equals its total distance to all other vertices.
Why it works
Every path through v either ends at v or has endpoints in different branches. Its length is the sum of the endpoint distances to v. For each endpoint in a branch of size s, exactly N−s possible partners lie outside that branch, including v. Hence d*(N−s) counts exactly its endpoint-distance contributions. Rerooting distance sums is correct because crossing an edge shortens distances to size[child] vertices and lengthens distances to the other N−size[child] vertices.
Watch out. Use 64-bit sums and products. Paths are unordered, but the branch formula adds each endpoint’s distance contribution, not the full path twice. Avoid recursive traversal on a 100000-vertex chain.
Commented C++17 solution
An independently written complete program with standard headers.
#include<iostream>
#include<vector>
// Algorithm: Root the tree.
using namespace std;
struct Edge{
int to;
int length;
};
int main(){
ios::sync_with_stdio(false);
cin.tie(nullptr);
int n;
cin>>n;
vector<vector<Edge>>graph(n);
for(int i=1;i<n;++i){
int a,b,length;
cin>>a>>b>>length;
--a;
--b;
graph[a].push_back({b,length});
graph[b].push_back({a,length});
}
vector<int>parent(n,-1),parentLength(n),order(1,0),size(n,1);
for(int i=0;i<static_cast<int>(order.size());++i){
int vertex=order[i];
for(Edge edge:graph[vertex]){
if(edge.to!=parent[vertex]){
parent[edge.to]=vertex;
parentLength[edge.to]=edge.length;
order.push_back(edge.to);
}
}
}
vector<long long>down(n,0),all(n,0);
for(int i=n-1;i>0;--i){
int vertex=order[i],above=parent[vertex];
size[above]+=size[vertex];
down[above]+=down[vertex]+1LL*parentLength[vertex]*size[vertex];
}
all[0]=down[0];
for(int i=1;i<n;++i){
int vertex=order[i];
all[vertex]=all[parent[vertex]]+1LL*parentLength[vertex]*(n-2*size[vertex]);
}
for(int vertex=0;vertex<n;++vertex){
long long answer=(all[vertex]-down[vertex])*size[vertex];
for(Edge edge:graph[vertex]){
if(parent[edge.to]==vertex){
long long distance=down[edge.to]+1LL*edge.length*size[edge.to];
answer+=distance*(n-size[edge.to]);
}
}
cout<<answer<<'\n';
}
}