NOIP2016提高组.天天爱跑步

题目

354. 天天爱跑步

算法标签: 树上倍增, $lca$, 树上差分, 贪心

思路

直接枚举每个点然后算路径对点的贡献是会超时的事件复杂度是$O(n^2)$级别, 但是如果枚举每条路径, 计算每个路径对每个点的贡献, 这样就能将时间复杂度压缩到$O(n)$级别, 为了处理对每个路径上点的贡献, 需要树上差分, 但是因为差分的数值是一组数, 不是一个数, 如果直接使用$map$存储, 也会超时, 因此也需要进行优化, 考虑对树上每个点的操作存储起来, 然后在$dfs$过程中计算贡献, 设对当前子树操作之前的数值是$val$, 那么最后答案需要累加上$s[u]-val$, $s[u]$是$dfs$子树后得到的每个点的贡献值

代码

#pragma GCC optimize(2)

#include <iostream>
#include <algorithm>
#include <cstring>
#include <vector>

using namespace std;

#define x first
#define y second

typedef pair<int, int> PII;
const int N = 300010, M = N << 1, K = 19;

int n, m;
int head[N], ed[M], ne[M], idx;
int w[N];
PII q[N];
int fa[N][K], depth[N];
vector<PII> op[N];
int ans[N], s[N * 3];

void add(int u, int v) {
   
	ed[idx] = v, ne[idx] = head[u], head[u] = idx++;
}

void dfs_fa(int u, int pre, int dep) {
   
	depth[u] = dep;
	for (int i = head[u]; ~i; i = ne[i]) {
   
		int v = ed[i];
		if (v == pre) continue;
		fa[v][0] = u;
		for (int k = 1; k < K; ++k) fa[v][k] = fa[fa[v][k - 1]][k - 1];
		dfs_fa(v, u, dep + 1);
	}
}

int lca(int u, int v) {
   
	if (depth[u] < depth[v]) swap(u, v);

	for (int k = K - 1; k >= 0; --k) {
   
		if (depth[fa[u][k]] >= depth[v]) {
   
			u = fa[u][k];
		}
	}

	if (u == v) return u;

	for (int k = K - 1; k >= 0; --k) {
   
		if (fa[u][k] != fa[v][k]) {
   
			u = fa[u][k];
			v = fa[v][k];
		}
	}

	return fa[u][0];
}

void dfs_sum(int u, int pre, int c) {
   
	int t = w[u] + c * depth[u] + N;
	int val = s[t];
	for (auto [x, y] : op[u]) s[x + N] += y;
	for (int i = head[u]; ~i; i = ne[i]) {
   
		int v = ed[i];
		if (v == pre) continue;
		dfs_sum(v, u, c);
	}
	ans[u] += s[t] - val;
}

// 计算节点u到lca
void calc_up() {
   
	for (int i = 0; i < m; ++i) {
   
		auto [u, v] = q[i];
		int p = lca(u, v);
		int t = depth[u];
		op[u].push_back({
   t, 1});
		op[fa[p][0]].push_back({
   t, -1});
	}
	dfs_sum(1, -1, 1);
}

// 计算lca到节点v
void calc_down() {
   
	for (int i = 0; i <= n; ++i) op[i].clear();
	memset(s, 0, sizeof s);
	for (int i = 0; i < m; ++i) {
   
		auto [u, v] = q[i];
		int p = lca(u, v);
		int t = depth[u] - 2 * depth[p];
		op[v].push_back({
   t, 1});
		op[p].push_back({
   t, -1});
	}
	dfs_sum(1, -1, -1);
}

int main() {
   
	ios::sync_with_stdio(false);
	cin.tie(0), cout.tie(0);

	memset(head, -1, sizeof head);

	cin >> n >> m;
	for (int i = 0; i < n - 1; ++i) {
   
		int u, v;
		cin >> u >> v;
		add(u, v), add(v, u);
	}

	for (int i = 1; i <= n; ++i) cin >> w[i];
	for (int i = 0; i < m; ++i) cin >> q[i].x >> q[i].y;

	dfs_fa(1, -1, 1);
	calc_up();
	calc_down();

	for (int i = 1; i <= n; ++i) cout << ans[i] << " ";
	cout << "\n";
	return 0;
}