P3478 STA-Station/换根 $dp$ 板子

P3478 [POI2008] STA-Station

link

给定一个 \(n\) 个点的树，请求出一个结点，使得以这个结点为根时，所有结点的深度之和最大。

一个结点的深度之定义为该节点到根的简单路径上边的数量。

对于全部的测试点，保证 \(1 \leq n \leq 10^6\)，\(1 \leq u, v \leq n\)，给出的是一棵树。

思路：树形dp（换根dp）。

我们设1节点为根节点，并先统计完1节点的答案。

接下来设计状态转移方程：

我们设 \(dp_{pos}\) 为以 \(pos\) 为根节点时，所有节点到根节点的深度之和。

设 \(to\) 为 \(pos\) 节点的子节点，则可以得到状态转移方程：

\[dp_{to} = dp_{pos} + (n-sz_{to}) - sz_{to} \]

其中 \(sz_{to}\) 是指 \(to\) 节点的子树的大小。

\((n-sz_{to})\) 是移到 \(to\) 节点后，除了 \(to\) 的子树的节点外的节点到 \(to\) 节点的距离增加的值。

\(- sz_{to}\) 则是移到 \(to\) 节点后，\(to\) 的子树的所有节点到 \(to\) 节点距离减少的值。

则这个过程可以用两遍 \(dfs\) 来解决。

第一遍：统计以 \(1\) 节点为根节点时候的答案，并预处理 \(sz\) 数组。

void predfs(int pos,int fa,int dep){
	dp[1] += dep;
	for(int to : tree[pos]){
		if(to != fa){
			predfs(to,pos,dep + 1);
			sz[pos] += sz[to];
		}
	}
}

第二遍：直接使用状态转移方程转移，注意是从根节点往子节点 \(dp\) 。

void dfs(int pos,int fa){
	for(int to : tree[pos]){
		if(to != fa){
			dp[to] = dp[pos] + n - 2 * sz[to];
			dfs(to,pos);
		}
	}
}

最后遍历整个 \(dp\) 数组并且输出答案即可。

\({\Huge code:}\)

#include<iostream>
#include<algorithm>
#include<vector>
using namespace std;
#define int long long
const int MAXN = 1e6 + 7;
vector<int> tree[MAXN];
int n;
int dp[MAXN];
int sz[MAXN];
void predfs(int pos,int fa,int dep){
	dp[1] += dep;
	for(int to : tree[pos]){
		if(to != fa){
			predfs(to,pos,dep + 1);
			sz[pos] += sz[to];
		}
	}
}
void dfs(int pos,int fa){
	for(int to : tree[pos]){
		if(to != fa){
			dp[to] = dp[pos] + n - 2 * sz[to];
			dfs(to,pos);
		}
	}
}
signed main(){
	ios::sync_with_stdio(false);
	cin.tie(0),cout.tie(0);
	cin>>n;
	for(int i = 1;i < n;i++){
		int x,y;
		cin>>x>>y;
		tree[x].push_back(y);
		tree[y].push_back(x);
	}
	for(int i = 1;i <= n;i++) sz[i] = 1;
	predfs(1,0,1);
	dfs(1,0);
	int ans = 0,id = 1;
	for(int i = 1;i <= n;i++){
		if(dp[i] > ans){
			ans = dp[i],id = i;
		}
	}
	cout<<id;
	return 0;
}

复杂度 \({\Large O(n)}\)，可以通过本题。