Colors on a Tree
National Olympiad — National round (2022 season)
Submission name: colors.cpp
Problem summary
A rooted tree supports recoloring a path, recoloring a subtree, and counting adjacent color changes on a path. The root is vertex 1.
Input
N; N−1 undirected edges; N initial colors; Q; then 1 U V C, 2 U C or 3 U V.
Output
One count for each type-3 query.
Constraints
1 ≤ N,Q ≤ 100000; vertices and colors are in [1,N].
Worked example
An independently prepared example for this explanation.
Input
6 1 2 2 6 2 3 3 4 3 5 1 1 3 1 2 1 5 3 6 4 1 6 4 2 3 1 5 2 2 4 3 3 5
Output
2 1 0
Solution analysis
- Use a heavy-light ordering which visits the heavy child first. A subtree is one contiguous interval, while a path splits into O(log N) heavy-chain intervals. Compute the decomposition iteratively so a chain of 100000 vertices cannot overflow the call stack.
- A segment summary stores its first color, last color, length and count of internal color changes. Concatenating A and B adds their counts and one exactly when A.last differs from B.first. Recoloring sets both boundary colors to the new color and the change count to zero, so a lazy assignment segment tree supports both interval operations.
- Path order matters: accumulate the upward pieces from U in a left summary, reversing their boundary colors. Prepend downward pieces from V to a right summary. Join both sides after the final common-chain piece. In the example 6→2→3→4 has colors 1,1,3,1 and therefore two changes. The path recoloring and subtree recoloring produce the later answers one and zero.
Why it works
Heavy-light decomposition makes every subtree contiguous and each path a logarithmic number of chain intervals because every light edge at least halves subtree size. The summary merge exactly counts internal boundaries and the single joining boundary; reversal exchanges endpoints without altering internal changes. Lazy assignment creates the exact constant-color summary. The two path accumulators preserve traversal order and include each path vertex once, so their final merge gives the correct answer.
Watch out. A merge is not commutative: boundary colors must follow the path direction. Include the LCA only once, and use iterative preprocessing on deep trees.
Commented C++17 solution
An independently written complete program with standard headers.
#include<iostream>
#include<vector>
#include<algorithm>
// Algorithm: Use a heavy-light ordering which visits the heavy child first.
using namespace std;
struct Summary{
int first=0,last=0,changes=0,length=0;
};
Summary join(Summary a,Summary b){
if(!a.length){
return b;
}
if(!b.length){
return a;
}
return {a.first,b.last,a.changes+b.changes+(a.last!=b.first),a.length+b.length};
}
Summary reverseSummary(Summary a){
swap(a.first,a.last);
return a;
}
struct Tree{
int n;
vector<Summary>data;
vector<int>lazy;
Tree(const vector<int>&a):n(static_cast<int>(a.size())),data(4*n),lazy(4*n){
build(1,0,n-1,a);
}
void build(int p,int l,int r,const vector<int>&a){
data[p].length=r-l+1;
if(l==r){
data[p].first=data[p].last=a[l];
return;
}
int m=(l+r)/2;
build(p*2,l,m,a);
build(p*2+1,m+1,r,a);
data[p]=join(data[p*2],data[p*2+1]);
}
void apply(int p,int color){
data[p].first=data[p].last=color;
data[p].changes=0;
lazy[p]=color;
}
Summary visit(int p,int l,int r,int a,int b,int color){
if(b<l||r<a){
return {};
}
if(a<=l&&r<=b){
if(color){
apply(p,color);
}
return data[p];
}
if(lazy[p]){
apply(p*2,lazy[p]);
apply(p*2+1,lazy[p]);
lazy[p]=0;
}
int m=(l+r)/2;
Summary result=join(visit(p*2,l,m,a,b,color),visit(p*2+1,m+1,r,a,b,color));
data[p]=join(data[p*2],data[p*2+1]);
return result;
}
Summary interval(int l,int r,int color=0){
return visit(1,0,n-1,l,r,color);
}
};
int main(){
ios::sync_with_stdio(false);
cin.tie(nullptr);
int n;
cin>>n;
vector<vector<int>>graph(n);
for(int i=1;i<n;++i){
int a,b;
cin>>a>>b;
--a;
--b;
graph[a].push_back(b);
graph[b].push_back(a);
}
vector<int>parent(n,-1),depth(n),size(n,1),heavy(n,-1),order{0};
for(int k=0;k<n;++k){
int v=order[k];
for(int u:graph[v]){
if(u!=parent[v]){
parent[u]=v;
depth[u]=depth[v]+1;
order.push_back(u);
}
}
}
for(int k=n-1;k>0;--k){
int v=order[k],p=parent[v];
size[p]+=size[v];
if(heavy[p]==-1||size[v]>size[heavy[p]]){
heavy[p]=v;
}
}
vector<int>head(n),position(n),stack{0};
int timer=0;
while(!stack.empty()){
int start=stack.back();
stack.pop_back();
for(int v=start;v!=-1;v=heavy[v]){
head[v]=start;
position[v]=timer++;
for(int u:graph[v]){
if(parent[u]==v&&u!=heavy[v]){
stack.push_back(u);
}
}
}
}
vector<int>colors(n);
for(int v=0;v<n;++v){
cin>>colors[position[v]];
}
Tree tree(colors);
int q;
cin>>q;
while(q--){
int type,u,v,color=0;
cin>>type>>u;
--u;
if(type==2){
cin>>color;
tree.interval(position[u],position[u]+size[u]-1,color);
continue;
}
cin>>v;
--v;
if(type==1){
cin>>color;
}
Summary left,right;
while(head[u]!=head[v]){
if(depth[head[u]]>=depth[head[v]]){
left=join(left,reverseSummary(tree.interval(position[head[u]],position[u],color)));
u=parent[head[u]];
}
else{
right=join(tree.interval(position[head[v]],position[v],color),right);
v=parent[head[v]];
}
}
if(depth[u]>=depth[v]){
left=join(left,reverseSummary(tree.interval(position[v],position[u],color)));
}
else{
left=join(left,tree.interval(position[u],position[v],color));
}
if(type==3){
cout<<join(left,right).changes<<'\n';
}
}
}