luogu P3629 [APIO2010]巡逻

#include <bits/stdc++.h>
using namespace std;
const int maxn = 100005;
const int maxm = 200005;
int he[maxn],ver[maxm],ne[maxm],cost[maxm],tot=1;
void add( int x,int y,int z ){
    ver[++tot] = y;
    ne[tot] = he[x];
    he[x] = tot;
    cost[tot] = z;
}
int son1[maxn],son2[maxn],h1[maxn],h2[maxn],res,s;
void dfs( int x,int f ){
    for( int cure = he[x];cure;cure = ne[cure] ){
        int y = ver[cure];
        if( y == f ) continue;
        dfs(y,x);
        int cur = h1[y]+cost[cure];
        if( h2[x] < cur ) {
            h2[x] = cur;
            son2[x] = cure;
            if (h2[x] > h1[x]) {
                swap(h1[x], h2[x]);
                swap(son1[x], son2[x]);
            }
        }
        if(res < h1[x]+h2[x])
        res = h1[x]+h2[x],s = x ;
    }
}
int main(){
    int n,k,x,y;
    scanf("%d%d",&n,&k);
    for( int i = 1;i < n;i++ ){
        scanf("%d%d",&x,&y);
        add(x,y,1);add(y,x,1);
    }
    memset( son1,0,sizeof(int)*(n+1) );
    memset( son2,0,sizeof(int)*(n+1) );
    memset( h1,0,sizeof(int)*(n+1) );
    memset(h2,0,sizeof(int)*(n+1));
    res = 0;
    dfs(1,0);
    x = s;
    int ans = (n-1)*2 -res+1;
    if( k == 1 ){
        printf("%d\n",ans);
        return 0;
    }
    for( int i = son1[x];i;i = son1[ver[i]] ){
        cost[i] = cost[i^1] = -1;
    }
    for( int i = son2[x];i;i = son1[ver[i]] ){
        cost[i] = cost[i^1] = -1;
    }
    memset( son1,0,sizeof(int)*(n+1) );
    memset( son2,0,sizeof(int)*(n+1) );
    memset( h1,0,sizeof(int)*(n+1) );
    memset(h2,0,sizeof(int)*(n+1));
    res = 0;
    dfs(1,0);
    ans = ans -res +1;
    printf("%d\n",ans);
    return 0;
}