Luogu P3258 松鼠的新家（树链剖分+线段树/树状数组）

题面

题解

　　这种题目一看就是重链剖分裸题，还是区间修改，单点查询，查询之前在遍历时要记一个$delta$，因为这一次的起点就是上一次的终点，不需要放糖，所以可以用$BIT$来写，但我写完$modify$才反应过来，所以没改了。

#include <cstdio>
#include <cstring>
#include <algorithm>
using std::swap;

const int N = 3e5 + 10;
int n, a[N], siz[N], son[N], fa[N], dep[N];
int tim, dfn[N], top[N], delta[N];
int cnt, from[N], to[N << 1], nxt[N << 1];
int val[N << 2], add[N << 2];
inline void addEdge(int u, int v) {
	to[++cnt] = v, nxt[cnt] = from[u], from[u] = cnt;
}

void dfs(int u) {
	siz[u] = 1, dep[u] = dep[fa[u]] + 1;
	for(int i = from[u]; i; i = nxt[i]) {
		int v = to[i]; if(v == fa[u]) continue;
		fa[v] = u, dfs(v), siz[u] += siz[v];
		if(siz[v] > siz[son[u]]) son[u] = v;
	}
}
void dfs(int u, int t) {
	dfn[u] = ++tim, top[u] = t;
	if(!son[u]) return ; dfs(son[u], t);
	for(int i = from[u]; i; i = nxt[i]) {
		int v = to[i];
		if(v != fa[u] && v != son[u]) dfs(v, v);
	}
}

inline void pushup (int o, int lc, int rc) { val[o] = val[lc] + val[rc]; }
inline void pushdown (int o, int lc, int rc, int len) {
	if(add[o]) {
		val[lc] += add[o] * (len - (len >> 1));
		val[rc] += add[o] * (len >> 1);
		add[lc] += add[o], add[rc] += add[o], add[o] = 0;
	}
}
void modify (int ml, int mr, int k, int o = 1, int l = 1, int r = n) {
	if(l >= ml && r <= mr) { val[o] += k * (r - l + 1), add[o] += k; return ; }
	int mid = (l + r) >> 1, lc = o << 1, rc = lc | 1; 
	pushdown(o, lc, rc, r - l + 1); 
	if(ml <= mid) modify(ml, mr, k, lc, l, mid);
	if(mr > mid) modify(ml, mr, k, rc, mid + 1, r);
	pushup(o, lc, rc);
}
int query(int qs, int o = 1, int l = 1, int r = n) {
	if(l == r && l == qs) return val[o];
	int mid = (l + r) >> 1, lc = o << 1, rc = lc | 1; 
	pushdown(o, lc, rc, r - l + 1); 
	if(qs <= mid) return query(qs, lc, l, mid);
	else return query(qs, rc, mid + 1, r);
}

inline void Path(int x, int y) {
	int fx = top[x], fy = top[y];
	while(fx != fy) {
		if(dep[fx] >= dep[fy]) modify(dfn[fx], dfn[x], 1), x = fa[fx], fx = top[x];
		else modify(dfn[fy], dfn[y], 1), y = fa[fy], fy = top[y];
	} if(dfn[x] > dfn[y]) swap(x, y); modify(dfn[x], dfn[y], 1);
}

int main () {
	scanf("%d", &n);
	for(int i = 1; i <= n; ++i) scanf("%d", a + i);
	for(int i = 1, u, v; i < n; ++i) {
		scanf("%d%d", &u, &v);
		addEdge(u, v), addEdge(v, u);
	}
	dfs(1), dfs(1, 1);
	for(int i = 2; i <= n; ++i)
		Path(a[i - 1], a[i]), ++delta[a[i]];
	for(int i = 1; i <= n; ++i)
		printf("%d\n", query(dfn[i]) - delta[i]);
	return 0;
}