dls的数据结构-习题课2

天天爱跑步

#include<bits/stdc++.h>

using namespace std;
const int N = 3e5+10;


int h[N], ne[2 * N], e[2 * N], idx;

void add(int a, int b){
	e[idx] = b, ne[idx] = h[a], h[a] = idx++;
}


// 注释部分可以维护信息，本代码执行为lca代码
const int LOGN = 21;
int depth[N], p[LOGN][N], n, m, W[N];
int l[N], r[N], tot;
// int val[LOGN][N];
// 求i==0的时候的值
void dfs1(int u, int fa){
	l[u] = ++ tot;
    depth[u] = depth[fa] + 1;
    for(int i = h[u]; i != -1; i = ne[i]){
        int j = e[i];
        if(j == fa) continue;
        dfs1(j, u);
        p[0][j] = u;
        // val[0][j] = w[i];
    }
    r[u] = tot;
}

void init(){
    // 这里注意val可能存在的边界问题
    for(int i = 1; i < LOGN; i ++){
        for(int j = 1; j <= n; j ++){
            p[i][j] = p[i - 1][p[i - 1][j]];
        }
    }
}

int query(int u, int v){
    // int res = 1 << 30;
    if(depth[u] < depth[v]) swap(u, v);
    int d = (depth[u] - depth[v]);
    for(int i = 0; i < LOGN; i ++){
        if(d >> i & 1){
            // res = min(res, val[i][u]);
            u = p[i][u];
        }
    }
    // if(u == v) return res;
    if(u == v) return u;

    for(int i = LOGN - 1; i >= 0; i --){
        if(p[i][u] != p[i][v]){
            // res = min(res, min(val[i][u], val[i][v]));
            u = p[i][u]; v = p[i][v];
        }
    }
    return p[0][u];
}

vector< array<int, 3> >up, down;
// s的大小要开两倍，因为up和down都是分成了两个
int ans[N], s[N * 2];

int main(){
	cin >> n >> m;
	memset(h, -1, sizeof h);
	for(int i = 1; i <= n - 1; i ++){
		int x, y; scanf("%d %d", &x, &y);
		add(x, y), add(y, x);
	}
	for(int i = 1; i <= n; i ++) scanf("%d", &W[i]);
	dfs1(1, 0); init();
	// 其实推一下式子发现，从起点的时候就能够确定这条路径上有没有符合条件的点了
	// 如果可以的话，我们就将答案存放在起点的位置
	// 另一个问题是，可能这条路径走不到根的位置然后他就会停下来
	// 所以我们在lca的位置向上带一个负数的标记
	for(int i = 1; i <= m; i ++){
		int u, v; scanf("%d %d", &u, &v);
		int w = query(u, v);
		up.push_back({depth[u], l[u], 1});
		up.push_back({depth[u], l[w], -1});
		int T = depth[u] + depth[v] - 2 * depth[w];
		down.push_back({T - depth[v], l[v], 1});
		down.push_back({T - depth[v], l[w], -1});
		// lca位置要单独判断
		if(W[w] == depth[u] - depth[w]) ans[w] ++;
	}
	m = up.size();
	sort(up.begin(), up.end());
	for(int i = 0; i < m; i ++){
		s[i + 1] = s[i] + up[i][2];
	}
	for(int v = 1; v <= n; v ++){
		int pl = lower_bound(up.begin(), up.end(), 
			array<int, 3> {W[v] + depth[v], l[v], -2}) - up.begin();
		int pr = lower_bound(up.begin(), up.end(),
			array<int, 3> {W[v] + depth[v], r[v], 2}) - up.begin();
		ans[v] += s[pr] - s[pl]; 
	}

	m = down.size();
	sort(down.begin(), down.end());
	for(int i = 0; i < m; i ++){
		s[i + 1] = s[i] + down[i][2];
	}
	for(int v = 1; v <= n; v ++){
		int pl = lower_bound(down.begin(), down.end(), 
			array<int, 3> {W[v] - depth[v], l[v], -2}) - down.begin();
		int pr = lower_bound(down.begin(), down.end(),
			array<int, 3> {W[v] - depth[v], r[v], 2}) - down.begin();
		ans[v] += s[pr] - s[pl]; 
	}
	for(int i = 1; i <= n; i ++) printf("%d%c", ans[i], " \n"[i == n]);

}