AIM Tech Round 3 (Div. 1) C. Centroids

Problem Description

给定一棵 \(n\) 个节点的树，你只能进行一次边替换。

边替换是指在树中删去一条边（不删除相关节点）并插入一条边（不添加新的节点）。

请你确定每个节点是否能在至多一次操作的情况下成为重心。

Input

第一行输入 \(n(2 \le n \le 400000)\) ，表示节点个数。

接下来 \(n-1\) ，每行输入两个整数 \(u_i,v_i(1 \le u_i,v_i \le n)\) ，表示 \(u_i\) 和 \(v_i\) 之间有一条边。

Output

输出 \(n\) 个整数，如果节点 \(i\) 能成为重心则输出 \(1\) ，否则输出 \(0\) 。

Solution

我们先从一个节点入手，进行分类讨论：

1. 所有子树重量都 \(\le \frac{n}{2}\) ，那么显然就是重心。

2. 只有一棵子树的重量 \(\le \frac{n}{2}\) ，由于我们只能删去一条边，那么我们必然是贪心地去删尽可能大的子树，然后直接与根节点相接，那么这棵尽可能大的子树重量应该同样满足 \(\le \frac{n}{2}\) ，那么就会分成以下两种情况：

   · \(weight_v-\max\limits_{weight_i \le \frac{n}{2}}(weight_i) \le \frac{n}{2}\) ，那么可以成为重心。

   · \(weight_v-\max\limits_{weight_i \le \frac{n}{2}}(weight_i) > \frac{n}{2}\) ，那么无法成为重心。

   上述 \(v\) 表示当前根节点 \(u\) 的儿子节点，\(i\) 表示 \(v\) 的儿子节点。

3. 如果重量 \(\le \frac{n}{2}\) 的子树大于等于 \(2\) ，那么显然无法成为重心。

如此一来，我们便发现要求一个节点是否能成为重心，就变成了求所有子树的重量以及所有子树中重量 \(\le \frac{n}{2}\) 的最大值。

对此，如果只要求一个节点，我们只需进行简单的树形 \(dp\) 即可。

令 \(dp_u\) 表示以 \(u\) 为父节点的子树重量中 \(\le \frac{n}{2}\) 的最大值，状态转移只需判断是否子树的 \(weight_v \le \frac{n}{2}\) ，如果成立 \(dp_u=\max(dp_u,weight_v)\) ，否则则为 \(dp_u=max(dp_u,dp_v)\) 。

如此一来我们便能得到某一个节点是否能成为重心，但是此题要求我们求出所有节点是否能成为重心，此时便通过换根 \(dp\) 来解决这个问题。

我们考虑如何将当前根节点 \(u\) 变成 \(v\) 的儿子节点， 设 \(dmax_0\) 和 \(dmax_1\) 分别表示以 \(u\) 为根节点时所有子树重量 \(\le \frac{n}{2}\) 的最大值和次大值。那么当 \(v\) 变成根时，只需判断其是否为 \(dmax_0\) 贡献子树，如果是则 \(dp_u=dmax_1\) ，否则 \(dp_u=dmax_0\) ，同时 \(weight_u=n-weight_v\)，然后每次访问一个新节点时判断其是否能成为重心即可。

最终复杂度为 \(O(n+m)\) 。

Code

#include <bits/stdc++.h>
#define endl '\n'
using namespace std;
const int N = 4e5 + 10, M = 8e5 + 10;
int n, ans[N];
int wet[N], dp[N];
int head[N], e[M], ne[M], idx = 0;
inline void addedge(int a, int b) {
	e[idx] = b, ne[idx] = head[a], head[a] = idx++;
}
void initial_dfs(int u, int fa) {
	wet[u] = 1;
	for (int i = head[u]; i != -1; i = ne[i]) {
		int v = e[i];
		if (v == fa)continue;
		initial_dfs(v, u);
		wet[u] += wet[v];
		if (wet[v] <= n / 2)dp[u] = max(dp[u], wet[v]);
		else dp[u] = max(dp[u], dp[v]);
	}
}
void dfs(int u, int fa) {
	ans[u] = 1;
	int cnt = 0;//统计u的所有子树中重量大于n/2的个数
	vector<int>dmax(2);
	for (int i = head[u]; i != -1; i = ne[i]) {
		int v = e[i];
		if (wet[v] > n / 2)cnt++;
		if (wet[v] > n / 2 && wet[v] - dp[v] > n / 2)ans[u] = 0;

		//统计x为以u为根节点时,所有子树weight<=n/2的最大值和次大值
		int x = wet[v] <= n / 2 ? wet[v] : dp[v];
		if (x >= dmax[0])dmax[1] = dmax[0], dmax[0] = x;
		else if (x > dmax[1])dmax[1] = x;
	}
	if (cnt >= 2)ans[u] = 0;

	for (int i = head[u]; i != -1; i = ne[i]) {
		int v = e[i];
		if (v == fa)continue;
		int x = wet[v] <= n / 2 ? wet[v] : dp[v];
		//判断最大值是否是当前转移的子树提供时
		if (x == dmax[0])dp[u] = dmax[1];
		else dp[u] = dmax[0];
		wet[u] = n - wet[v];
		dfs(v, u);
	}
}
signed main() {
	ios::sync_with_stdio(false);
	cin.tie(0);
	cout.tie(0);
	memset(head, -1, sizeof head);
	cin >> n;
	for (int i = 1; i <= n - 1; i++) {
		int u, v;
		cin >> u >> v;
		addedge(u, v), addedge(v, u);
	}
	initial_dfs(1, -1);
	dfs(1, -1);
	for (int i = 1; i <= n; i++)cout << ans[i] << ' ';
	cout << endl;
}