题面
直接写正解吧, 不想再写部分分了, 对于\(u\)和\(v\), 我们可以将它拆成两条路径, \(u\)到\(lca(u, v)\)和\(lca(u, v)\)到v, 在这里只分析从\(u\)到\(lca(u, v)\)的路径(其实是我不想写).
对于一个点\(u\), 设它的深度为\(dep[u]\), 值为\(dep[i] + w[i]\)的物品有\(cnt[w[i] + dep[i]]\)个, 这样的话, 只有在满足\(dep[u] - dep[i] = w[i]\)的时候, 才能对这个点产生贡献, 由于这三个数都是确定的, 所以在某个起点出发到一个终点结束的路径, 我们处理出他的lca后, 可以使用树上差分, 将一个值为\(dep[u]\)的物品从\(u\)传入, 从\(lca(u, v)\)的父亲传出, 这样就可以对每个点进行树上差分, 在插入之前记录一下\(cnt[w[i] + dep[i]]\)的物品有多少个, 再从底往上更新, 查询当前点\(dep[i] + w[i]\)的物品有多少个, 相减即可.
至于\(lca(u, v)\)的情况, 大家根据上面列一个式子用同样的方法判断就行了, 这里只列出式子, 大家也可以自己推导一下, 最后是这个样子\(dep[u] + dep[i] - 2 * dep[lca(u, v)] = w[i]\), 有可能等式左边的会小于零, 在数组上平移一段即可.
具体代码
#include <iostream>
#include <cstdio>
#include <vector>
#define N 300005
using namespace std;
int n, m, w[N], head[N], cnt, f[N][20], dep[N], num[N], c1[N], c2[N << 1];
struct node
{
int to, next;
} edge[N << 1];
vector<int> lin[N], lout[N], rin[N], rout[N];
inline int read()
{
int x = 0, w = 1;
char c = getchar();
while(c < '0' || c > '9') { if (c == '-') w = -1; c = getchar(); }
while(c >= '0' && c <= '9') { x = x * 10 + c - '0'; c = getchar(); }
return x * w;
}
inline void add(int u, int v) { edge[++cnt].to = v; edge[cnt].next = head[u]; head[u] = cnt; }
void dfs(int u, int fa)
{
dep[u] = dep[fa] + 1;
for(int i = 1; i <= 18 && f[f[u][i - 1]][i - 1]; i++)
f[u][i] = f[f[u][i - 1]][i - 1];
for(int i = head[u]; i; i = edge[i].next)
{
int v = edge[i].to; if(v == fa) continue;
f[v][0] = u; dfs(v, u);
}
}
int LCA(int u, int v)
{
if(dep[u] > dep[v]) swap(u, v);
for(int i = 18; i >= 0; i--)
if(dep[f[v][i]] >= dep[u]) v = f[v][i];
if(v == u) return u;
for(int i = 18; i >= 0; i--)
if(f[u][i] != f[v][i]) { u = f[u][i]; v = f[v][i]; }
return f[u][0];
}
void down(int u)
{
int sum1 = c1[w[u] + dep[u]], sum2 = c2[w[u] - dep[u] + 300000];
for(int i = head[u]; i; i = edge[i].next)
if(edge[i].to != f[u][0]) down(edge[i].to);
for(unsigned int i = 0; i < lin[u].size(); i++) c1[lin[u][i]]++;//记录一下
for(unsigned int i = 0; i < lout[u].size(); i++) c1[lout[u][i]]--;
for(unsigned int i = 0; i < rin[u].size(); i++) c2[rin[u][i] + 300000]++;
for(unsigned int i = 0; i < rout[u].size(); i++) c2[rout[u][i] + 300000]--;//记得平移
num[u] += c1[w[u] + dep[u]] - sum1 + c2[w[u] - dep[u] + 300000] - sum2; //相减即为答案
}
int main()
{
n = read(); m = read();
for(int i = 1; i < n; i++)
{
int u = read(), v = read();
add(u, v); add(v, u);
}
for(int i = 1; i <= n; i++) w[i] = read();
dfs(1, 0);
for(int i = 1; i <= m; i++)
{
int u = read(), v = read();
int lca = LCA(u, v);
lin[u].push_back(dep[u]); lout[f[lca][0]].push_back(dep[u]); rin[v].push_back(dep[u] - 2 * dep[lca]); rout[lca].push_back(dep[u] - 2 * dep[lca]); //u到lca和lca到v, 由于lca已经走过, 这个时候从lca传入就相当于这个物品是在lca与v那条链上lca的儿子传入, 差分嘛
}
down(1);
for(int i = 1; i <= n; i++) printf("%d%c", num[i], i == n ? '\n' : ' ');
return 0;
}