Codeforces 709E 树的重心

题意

传送门 Codeforces 709E Centroids

题解

DFS 求出置换边 $0$ 次情况下，所有的树的重心。任取一个重心再次 DFS，此时任意子树的规模不大于 $\lfloor n/2\rfloor$。那么对于任意子树 $i$ 上的节点 $v$，若其在置换边 $1$ 次的情况下为树的重心，则删除的边有两种情况：子树 $i$ 与根节点的连边，将根节点连接到 $v$ 即可；根节点与其余子节点的连边（根据贪心策略，任意子树已经满足规模不大于 $\lfloor n/2\rfloor$，此时树的其余部分的节点数尽可能的少），需要判断 $v$ 是否满足条件。

#include <bits/stdc++.h>
using namespace std;
#define pb push_back
typedef long long ll;
const int MAXN = 4E5 + 5;
int N, lim, sz[MAXN], res[MAXN];
vector<int> G[MAXN];
map<int, int> cnt;

void dfs(int u, int p, int &rt, int &mx)
{
    sz[u] = 1;
    int t = -1;
    for (auto &v : G[u])
        if (v != p)
        {
            dfs(v, u, rt, mx);
            sz[u] += sz[v];
            t = max(t, sz[v]);
        }
    t = max(t, N - sz[u]);
    if (t <= lim)
        res[u] = 1;
    if (mx == -1 || t < mx)
        rt = u, mx = t;
}

void get_sz(int u, int p)
{
    sz[u] = 1;
    for (int v : G[u])
        if (v != p)
            get_sz(v, u), sz[u] += sz[v];
}

void get_res(int u, int p)
{
    res[u] = 1;
    for (auto &v : G[u])
        if (v != p)
            get_res(v, u);
}

void find(int u, int p)
{
    int s = 1;
    for (auto &v : G[u])
        if (v != p)
            s += sz[v];
    if (cnt.lower_bound(N - lim - s) != cnt.end())
        res[u] = 1;
    for (auto &v : G[u])
        if (v != p)
            find(v, u);
}

int main()
{
    ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
    cin >> N;
    lim = N / 2;
    for (int i = 1; i < N; ++i)
    {
        int u, v;
        cin >> u >> v;
        --u, --v;
        G[u].pb(v), G[v].pb(u);
    }
    int rt = -1, mx = -1;
    dfs(0, -1, rt, mx);
    get_sz(rt, -1);
    for (auto &u : G[rt])
        ++cnt[sz[u]];
    for (auto &u : G[rt])
    {
        if (N - sz[u] <= lim)
            get_res(u, rt);
        if (--cnt[sz[u]] == 0)
            cnt.erase(sz[u]);
        find(u, rt);
        ++cnt[sz[u]];
    }
    for (int i = 0; i < N; ++i)
        cout << res[i] << (i + 1 == N ? '\n' : ' ');
    return 0;
}