P3478 [POI2008] STA-Station

题目链接：

既然让求深度之和，那么我就定义以 \(i\) 为根时深度之和为 \(f_i\)，现在就是思考状态转移的问题。如果以某种手段得到了 \(f_1\)，那么接下来的转移就好说了。
设 \(u\) 为当前节点，\(j\) 是当前节点的子节点。\(s_i\) 表示以 \(i\) 为根的子树中的节点数量，则 \(s_u=1+ \sum{s_j}\)。当以 \(u\) 为根转到以 \(j\) 为根时，

• 所有在 \(j\) 的子树上的节点深度都减少了 \(1\)，那么总深度和就减少了 \(s_j\)
• 所有不在 \(j\) 的子树上的节点深度都增加了 \(1\)，那么总深度和就增加了 \(n - s_j\)

由此可以得出状态转移方程：\(f_j = f_u - s_j + n - s_j = f_u + n - 2 \cdot s_j\)

因此得出思路：

1. 先进行一次 \(\rm DFS\)，求出以 \(1\) 为根时每个点的深度以及以某个节点为根时其子树中的节点总数 \(s_i\)
2. 计算出 \(1\) 号节点的答案 \(f[1]\)
3. 再进行一次 \(\rm DFS\)，递推求出每个点的答案

#include <bits/stdc++.h>

using namespace std;
using LL = long long;

const int N = 1e6 + 5, M = N << 1;
int e[M], h[N], ne[M];
int n, idx;
LL s[N], dep[N], f[N];//f[i]表示以i为根时所有节点的深度之和，dep[i]表示以1为根时节点i的深度，s[i]表示以i为根的子树中的节点数量 

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

void dfs(int u, int father) {
	s[u] = 1;//根节点也算一个
	dep[u] = dep[father] + 1;
	for (int i = h[u]; i != -1; i = ne[i]) {
		int j = e[i];
		if (j == father) continue;
		dfs(j, u);
		s[u] += s[j];//加上所有子树的节点
	} 
}

void get_ans(int u, int father) {
	for (int i = h[u]; i != -1; i = ne[i]) {
		int j = e[i];
		if (j == father) continue;
		f[j] = f[u] - 2 * s[j] + n;
		get_ans(j, u);
	}
}

int main()
{
	ios::sync_with_stdio(false), cin.tie(nullptr);
	cin >> n;
	memset(h, -1, sizeof h);
	for (int i = 1; i < n; i++) {
		int a, b;
		cin >> a >> b;
		add(a, b), add(b, a);
	}
	dfs(1, -1);
	for (int i = 1; i <= n; i++) f[1] += dep[i]; 
	get_ans(1, -1); 
	
	LL ans = -1e9, id = 1;
	for (int i = 1; i <= n; i++) {
		if (f[i] > ans) {
			ans = f[i];
			id = i;
		}
	}
	cout << id;
	return 0; 
}

注：
① 转移方程题目问什么设什么先，求不出来再考虑换或者辅助一下 \(dp\)
② \(dp\) 数组要开 \(\rm long\) \(\rm long\)
③ 求 \(f_1\) 的时候只能传两个参数，所以 \(\rm dep\) 需要用数组记录

引申：
我们其实根本不需要求出每个点的深度，也即 \(\rm dep\) 数组可以不用。
通过简单证明可以发现，若定义根节点的深度为 \(1\)，则所有点的深度之和等于所有点的子树大小之和。
我们可以从贡献的角度考虑：
对于一个深度为 \(k\) 的节点，它会给从它到根节点路径上所有点的子树大小提供 \(1\) 的贡献，而这条路径上的点有 \(k\) 个，证毕。

#include <bits/stdc++.h>

using namespace std;
using LL = long long;

const int N = 1e6 + 5, M = N << 1;
int e[M], h[N], ne[M];
int n, idx;
LL s[N], f[N];//f[i]表示以i为根时所有节点的深度之和，dep[i]表示以1为根时节点i的深度，s[i]表示以i为根的子树中的节点数量 

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

void dfs(int u, int father) {
	s[u] = 1;
	for (int i = h[u]; i != -1; i = ne[i]) {
		int j = e[i];
		if (j == father) continue;
		dfs(j, u);
		s[u] += s[j];
	} 
}

void get_ans(int u, int father) {
	for (int i = h[u]; i != -1; i = ne[i]) {
		int j = e[i];
		if (j == father) continue;
		f[j] = f[u] - 2 * s[j] + n;
		get_ans(j, u);
	}
}

int main()
{
	ios::sync_with_stdio(false), cin.tie(nullptr);
	cin >> n;
	memset(h, -1, sizeof h);
	for (int i = 1; i < n; i++) {
		int a, b;
		cin >> a >> b;
		add(a, b), add(b, a);
	}
	dfs(1, -1);
	for (int i = 1; i <= n; i++) f[1] += s[i]; 
	get_ans(1, -1); 
	
	LL ans = -1e9, id = 1;
	for (int i = 1; i <= n; i++) {
		if (f[i] > ans) {
			ans = f[i];
			id = i;
		}
	}
	cout << id;
	return 0; 
}