98 lines
2.6 KiB
C++
98 lines
2.6 KiB
C++
#include <bits/stdc++.h>
|
||
using namespace std;
|
||
const int N = 600010, M = N << 1;
|
||
int n, m;
|
||
int a[N];
|
||
|
||
//链式前向星
|
||
int e[M], h[N], idx, w[M], ne[M];
|
||
void add(int a, int b, int c = 0) {
|
||
e[idx] = b, ne[idx] = h[a], w[idx] = c, h[a] = idx++;
|
||
}
|
||
|
||
//树上差分模板题【点权】
|
||
int depth[N], f[N][31];
|
||
int dlt[N]; //差分数组
|
||
|
||
// 倍增2^k,计算k的办法
|
||
// T = log(n) / log(2) + 1;
|
||
const int T = 25;
|
||
|
||
//倍增求 a,b的最近公共祖先
|
||
void bfs(int root) {
|
||
queue<int> q;
|
||
q.push(root);
|
||
depth[root] = 1;
|
||
|
||
while (q.size()) {
|
||
int u = q.front();
|
||
q.pop();
|
||
for (int i = h[u]; ~i; i = ne[i]) {
|
||
int j = e[i];
|
||
if (depth[j]) continue;
|
||
depth[j] = depth[u] + 1;
|
||
f[j][0] = u;
|
||
for (int k = 1; k <= T; k++) f[j][k] = f[f[j][k - 1]][k - 1];
|
||
q.push(j);
|
||
}
|
||
}
|
||
}
|
||
int lca(int a, int b) {
|
||
if (depth[a] < depth[b]) swap(a, b);
|
||
for (int k = T; k >= 0; k--)
|
||
if (depth[f[a][k]] >= depth[b])
|
||
a = f[a][k];
|
||
|
||
if (a == b) return a;
|
||
for (int k = T; k >= 0; k--)
|
||
if (f[a][k] != f[b][k])
|
||
a = f[a][k], b = f[b][k];
|
||
|
||
return f[a][0];
|
||
}
|
||
|
||
//前缀和
|
||
void dfs(int u, int fa) {
|
||
for (int i = h[u]; ~i; i = ne[i]) {
|
||
int j = e[i];
|
||
if (j == fa) continue;
|
||
dfs(j, u);
|
||
dlt[u] += dlt[j];
|
||
}
|
||
}
|
||
int main() {
|
||
memset(h, -1, sizeof h);
|
||
scanf("%d", &n);
|
||
|
||
for (int i = 1; i <= n; i++) scanf("%d", &a[i]); //需要按a[i]规定的路线逐个走
|
||
|
||
//表示标号 a 和 b 的两个房间之间有树枝相连
|
||
for (int i = 1; i < n; i++) {
|
||
int a, b;
|
||
scanf("%d %d", &a, &b);
|
||
add(a, b), add(b, a);
|
||
}
|
||
//预处理出lca的depth数组和f倍增数组
|
||
bfs(1);
|
||
|
||
// lca查表+树上点权差分
|
||
for (int i = 1; i < n; i++) { // n-1条边
|
||
int x = a[i], y = a[i + 1];
|
||
int lc = lca(x, y);
|
||
//点权
|
||
dlt[x]++;
|
||
dlt[y]++;
|
||
dlt[lc]--;
|
||
dlt[f[lc][0]]--;
|
||
}
|
||
//将差分还原回原始数组
|
||
dfs(1, 0);
|
||
|
||
//我们仔细想想可以发现，每一个路径的终点又是下条路径的起点，而我们对其修改了两遍，所以这个算法就有问题了
|
||
//对每条路径修改后，将终点的值减1：这样的话，就不存在重复覆盖的问题了，而且这也恰好符合了终点要 -1的情况
|
||
for (int i = 2; i <= n; i++) dlt[a[i]]--;
|
||
|
||
//输出每个房间需要放的糖果数量
|
||
for (int i = 1; i <= n; i++) printf("%d\n", dlt[i]);
|
||
return 0;
|
||
} |