洛谷-P3478 STA-Station

STA-Station

换根dp 模板

去到相邻的点可以根据去到的点的子树有多少个结点，来调整当前的值

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <vector>
#include <queue>
#include <array>
using namespace std;
typedef long long ll;
vector<vector<int>>gra;
vector<int>siz;
int n;

ll dfs(int now, int pre)
{
    siz[now] = 1;
    int ans = 1;
    for(int nex : gra[now])
    {
        if(nex == pre) continue;
        ans += dfs(nex, now);
        siz[now] += siz[nex];
    }
    return ans;
}

int bfs(int s, ll sum)
{
    queue<array<ll, 2>>q;
    vector<int>vis(n + 1), tot(n + 1, 0);
    q.push({s, sum});
    while(q.size())
    {
        auto [now, val] = q.front();
        q.pop();
        if(vis[now]) continue;
        vis[now] = 1;
        tot[now] = val;
        if(val > tot[s]) s = now;
        for(int nex : gra[now])
        {
            if(vis[nex]) continue;
            q.push({nex, val - siz[nex] - siz[nex] + n});
        }
    }
    return s;
}

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
    cin >> n;
    gra.resize(n + 1);
    siz.resize(n + 1);
    for(int i=1; i<n; i++)
    {
        int a, b;
        cin >> a >> b;
        gra[a].push_back(b);
        gra[b].push_back(a);
    }
    ll sum = dfs(1, 1);
    int rt = bfs(1, sum);
    cout << rt << endl;
    return 0;
}