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}$ ，那么就会分成以下两种情况：

   · $weigh{t}_v-\underset{weigh{t}_i\le \frac{n}{2}}{\max }(weigh{t}_i)\le \frac{n}{2}$ ，那么可以成为重心。

   · $weigh{t}_v-\underset{weigh{t}_i\le \frac{n}{2}}{\max }(weigh{t}_i)>\frac{n}{2}$ ，那么无法成为重心。

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

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

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

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

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

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

我们考虑如何将当前根节点 $u$ 变成 $v$ 的儿子节点， 设 $dma{x}_0$ 和 $dma{x}_1$ 分别表示以 $u$ 为根节点时所有子树重量 $\le \frac{n}{2}$ 的最大值和次大值。那么当 $v$ 变成根时，只需判断其是否为 $dma{x}_0$ 贡献子树，如果是则 $d{p}_u=dma{x}_1$ ，否则 $d{p}_u=dma{x}_0$ ，同时 $weigh{t}_u=n-weigh{t}_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;
}