All Subtree Investigations
National Olympiad — Regional round (2025 season)
Submission name: tree.cpp
Problem summary
A tree is rooted at vertex 1. For every label interval [l,r], mark every vertex whose ancestor, including itself, has a label in this interval. Find the sum of the numbers of marked vertices over all 1≤l≤r≤n.
Input
n, followed by the parents of vertices 2 through n.
Output
The exact sum over all intervals.
Constraints
1 ≤ n ≤ 10⁵; parent labels are in [1,n] and need not be smaller than their children; the input is a tree.
Worked example
An independently prepared example for this explanation.
Input
6 5 5 1 1 4
Output
87
Solution analysis
- Reverse the count: for each vertex v, count label intervals that contain at least one ancestor of v. There are total=n(n+1)/2 intervals. If the ancestor labels are sorted, an interval misses all of them exactly when it lies entirely inside one gap between consecutive labels, including the ends 0 and n+1.
- A gap containing g absent labels contributes g(g+1)/2 missed intervals. During DFS maintain an ordered set of labels on the root-to-current path and the sum missed over its gaps. Inserting a label splits one gap into two; subtract the old triangular count and add the two new ones. Undo that operation on leaving the vertex.
- Add total−missed for each vertex after inserting its own label. In the first example, ancestors of vertex 2 have labels {1,5,2}. Sorted, these leave a gap {3,4} and gap {6}, so four intervals miss them. Of all 21 intervals, 17 mark vertex 2. The same calculation for every vertex totals 87. Use an explicit event stack to handle a path of 100,000 vertices.
Why it works
A vertex is marked by [l,r] if and only if that interval contains one of its ancestor labels. Every interval that avoids the labels belongs to one unique gap, so the triangular-gap formula counts the complement exactly. The set update replaces precisely the affected gap and preserves missed. DFS insertion and removal make the set equal to the current vertex’s ancestor path; summing its complement count therefore counts each marked vertex in each investigation exactly once.
Watch out. Ancestor means the path to the root, not smaller labels. Include the vertex itself. Restore the set when returning from a subtree, and use 64-bit counts.
Commented C++17 solution
An independently written complete program with standard headers.
#include<iostream>
#include<iterator>
#include<set>
#include<vector>
using namespace std;
long long intervals(long long size){
return size*(size+1)/2;
}
int main(){
ios::sync_with_stdio(false);
cin.tie(nullptr);
int n;
cin>>n;
vector<vector<int>>children(n+1);
for(int vertex=2;vertex<=n;++vertex){
int parent;
cin>>parent;
children[parent].push_back(vertex);
}
set<int>ancestors{0,n+1};
long long total=intervals(n);
long long missed=total;
long long answer=0;
vector<int>events{1};
while(!events.empty()){
int event=events.back();
events.pop_back();
int vertex=event>0?event:-event;
if(event>0){
auto right=ancestors.lower_bound(vertex);
int next=*right;
int previous=*prev(right);
missed-=intervals(next-previous-1);
missed+=intervals(vertex-previous-1)+intervals(next-vertex-1);
ancestors.insert(vertex);
answer+=total-missed;
events.push_back(-vertex);
for(int child:children[vertex]){
events.push_back(child);
}
}
else{
auto current=ancestors.find(vertex);
int previous=*prev(current);
int next=*std::next(current);
// Restore the unsplit gap before visiting a sibling subtree.
missed-=intervals(vertex-previous-1)+intervals(next-vertex-1);
missed+=intervals(next-previous-1);
ancestors.erase(current);
}
}
cout<<answer<<'\n';
}