【Gym 102059A】Coloring Roads（树链剖分+单调栈）

https://codeforc.es/gym/102059/problem/A

题意

给出一棵树。
询问u, c, m: 将结点u到根节点路径上的边都染色成c，询问染色边数为m的颜色有多少种。

题解

每次都是从一个结点到根，所以可以对这颗树进行树链剖分(heavy-light decompose)

维护
has[col]代表颜色col有多少条边。
cnt[m]边数为m的颜色有多少个。
dfn[v]:结点v的dfs编号。
对于每条链的头维护一个dfn的单调递增栈。
栈中的内容为(dfn,c)，从dfn处用颜色c染色到链头。

对于一个染色操作，模拟树链跑到根结点的过程。

如果当前的结点dfn大于栈头的dfn，那就要将栈头的这个颜色的影响去掉，并弹出。

如果当前的结点的dfn小于栈头的dfn，那就将当前dfn到链头的之前颜色的影响去掉。

将dfn到链头染色为c，然后加入单调栈中。

代码

官方std

#include 
using namespace std;
using lint = long long;
using pi = pair<int, int>;
const int MAXN = 200005;

int n, c, q, dfn[MAXN], chn[MAXN], sz[MAXN], par[MAXN], piv;
vector<int> gph[MAXN];

void dfs0(int x){
	sz[x] = 1;
	for(auto &i : gph[x]){
		gph[i].erase(find(gph[i].begin(), gph[i].end(), x));
		par[i] = x;
		dfs0(i);
		sz[x] += sz[i];
	}
}

void dfs1(int x){
	dfn[x] = ++piv;
	if(gph[x].empty()) return;
	chn[gph[x][0]] = chn[x];
	dfs1(gph[x][0]);
	for(int i=1; i<gph[x].size(); i++){
		chn[gph[x][i]] = gph[x][i];
		dfs1(gph[x][i]);
	}
}

int has[MAXN], cnt[MAXN];
vector<pi> stk[MAXN];

void change_color(int col, int c){
	cnt[has[col]]--;
	has[col] += c;
	cnt[has[col]]++;
}

int main(){
	scanf("%d %d %d",&n,&c,&q);
	cnt[0] = c;
	for(int i=0; i<n-1; i++){
		int s, e;
		scanf("%d %d",&s,&e);
		gph[s].emplace_back(e);
		gph[e].emplace_back(s);
	}
	dfs0(1);
	chn[1] = 1;
	for(int i=1; i<=n; i++){
		sort(gph[i].begin(), gph[i].end(), [&](const int &a, const int &b){
			return sz[a] > sz[b]; 
		});
	}
	dfs1(1);
	while(q--){
		int v, c;
		scanf("%d %d",&v,&c);
		while(v){
		    int ch = chn[v];
			int prv = dfn[ch] - 1;
			while(!stk[ch].empty() && stk[ch].back().first <= dfn[v]){
				auto val = stk[ch].back();
				stk[ch].pop_back();
				change_color(val.second, -(val.first - prv));
				prv = val.first;
			}
			if(!stk[ch].empty()){
				auto val = stk[ch].back();
				change_color(val.second, -(dfn[v] - prv));
			}
			change_color(c, dfn[v] - dfn[ch] + 1);
			stk[ch].emplace_back(dfn[v], c);
			v = par[ch];
		}
		int m;
		scanf("%d",&m);
		printf("%d\n", cnt[m] - (has[c] == m) + (has[c] == m + 1));
	}
}