【BZOJ-1912】patrol巡逻 树的直径 + DFS（树形DP）

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cmath>
using namespace std;
int read()
{
    int x=0,f=1;char ch=getchar();
    while(ch<‘0‘||ch>‘9‘){if(ch==‘-‘)f=-1;ch=getchar();}
    while(ch>=‘0‘&&ch<=‘9‘){x=x*10+ch-‘0‘;ch=getchar();}
    return x*f;
}
#define maxn 100010
int n,k,ans,zj,s;
struct EdgeNode{int next,to,len;}edge[maxn<<1];
int head[maxn],cnt=1;int road[maxn]={0},croad[maxn]={0};
void add(int u,int v,int w)
{
    cnt++;
    edge[cnt].to=v; edge[cnt].next=head[u]; head[u]=cnt; edge[cnt].len=w;
}
void insert(int u,int v,int w) {add(u,v,w); add(v,u,w);}
int DFS(int now,int fa)
{
    int maxd=0,cmaxd=0;
    for (int i=head[now]; i; i=edge[i].next)
        if (edge[i].to!=fa)
            {
                int len=DFS(edge[i].to,now)+edge[i].len; 
                if (len>maxd) cmaxd=maxd,maxd=len,croad[now]=road[now],road[now]=i;            
                else if (len>cmaxd) cmaxd=len,croad[now]=i;
            }
    if (maxd+cmaxd>zj) zj=maxd+cmaxd,s=now;
//    printf("%d %d %d %d\n",now,fa,maxd,cmaxd);
    return maxd;
}
int main()
{
    n=read(); k=read();
    for (int u,v,i=1; i<=n-1; i++) u=read(),v=read(),insert(u,v,1);
    DFS(1,0); ans=2*(n-1)-zj+1;
    if (k==2)
        {
            for (int i=road[s]; i; i=road[edge[i].to]) edge[i].len=edge[i^1].len=-1;
            for (int i=croad[s]; i; i=road[edge[i].to]) edge[i].len=edge[i^1].len=-1;
            zj=0; DFS(1,0); ans=ans-zj+1;
        }
    printf("%d\n",ans);
    return 0;
}