loj10159. 「一本通 5.2 练习 2」旅游规划

#include<cstdio>
#include<iostream>
#include<cstdlib>
#include<vector>
using namespace std;
const int maxn = 200010;
void qread(int &x) {
    x = 0;
    register int ch = getchar();
    while(ch < ‘0‘ || ch > ‘9‘)    ch = getchar();
    while(ch >= ‘0‘ && ch <= ‘9‘)    x = 10 * x + ch - 48, ch = getchar();
}
int n, rt = 0;

int deep[maxn];
int f[maxn];

int head[maxn];
int go[maxn << 1];
int nxt[maxn << 1];

int d1[maxn];
int d2[maxn];
vector <int> c1[maxn];
vector <int> c2[maxn];
vector <int> v;
int sign[maxn];
void dfs(int x){
    for(register int i = head[x]; i; i = nxt[i]){
        if(!deep[go[i]]){
            f[go[i]] = x;
            deep[go[i]] = deep[x] + 1;
            dfs(go[i]);
        }
    }
}
inline void init(){
    qread(n);
    for(int i=1; i<n; ++i){
        int x, y;
        qread(x), qread(y);
        go[i] = y;
        nxt[i] = head[x];
        head[x] = i;
        go[i + n] = x;
        nxt[i + n] = head[y];
        head[y] = i + n;
    }
    deep[rt] = 1;
    dfs(rt);
}
void DP(int x){
    for(register int i = head[x]; i; i =nxt[i]){
        if(go[i] != f[x]){
            DP(go[i]);
            if(d1[go[i]] + 1 > d1[x]){
                d2[x] = d1[x];
                d1[x] = d1[go[i]] + 1;
                c2[x].clear();
                for(int j=0; j<c1[x].size(); ++j)
                    c2[x].push_back(c1[x][j]);
                c1[x].clear();
                c1[x].push_back(go[i]);
            }
            else if(d1[go[i]] + 1 == d1[x])
                c1[x].push_back(go[i]);
            else if(d1[go[i]] + 1 > d2[x]){
                d2[x] = d1[go[i]] + 1;
                c2[x].clear();
                c2[x].push_back(go[i]);
            }
            else if(d1[go[i]] + 1 == d2[x]){
                c2[x].push_back(go[i]);
            }
        }
    }
}
void dfs2(int x){
    if(!x)    return ;
    sign[x] = 1;
    for(int i=0; i<c1[x].size(); ++i)
        dfs2(c1[x][i]);
}
int main(void){
    init();
    DP(rt);
    int p = 0;
    for(int i=0; i<n; ++i){
        int np = c1[i].size() > 1 ? d1[i] << 1 : d1[i] + d2[i];
        p = max(p, np);
    }
    for(int i=0; i<n; ++i){
        int np = c1[i].size() > 1 ? d1[i] << 1 : d1[i] + d2[i];
        if(np == p)    v.push_back(i);
    }
    for(int i=0; i<v.size(); ++i){
        rt = v[i];
        sign[rt] = 1;
        if(c1[rt].size()>=2){
            for(int i=0; i<c1[rt].size(); ++i)
                dfs2(c1[rt][i]);
        }
        else{
            dfs2(c1[rt][0]);
            for(int i=0; i<c2[rt].size(); ++i)
                dfs2(c2[rt][i]);
        }    
    }
    for(int i=0; i<n; ++i)
        if(sign[i])
            cout << i << endl;
}