bzoj 1912 [Apio2010]patrol 巡逻 树形dp

如果K=1那么就加连接直径两段的边。
否则把直径上的边权赋为-1，然后再求一遍直径。
答案是(n-1)*2-(两次求出的直径和)+K

#include <bits/stdc++.h>
using namespace std;
#define N 110000
int n,K,tot;
int f[N],head[N],nex[N<<1],to[N<<1];
int bj[N],deep[N],fa[N],pos[N];
int p1,p2,ans,fin;
void add(int x,int y)
{
    tot++;
    nex[tot]=head[x];head[x]=tot;
    to[tot]=y;
}
void dfs(int x,int y)
{
    deep[x]=deep[y]+1;
    f[x]=0;pos[x]=x;fa[x]=y;
    for(int i=head[x];i;i=nex[i])
        if(to[i]!=y)
        {
            dfs(to[i],x);
            if(f[x]+f[to[i]]+bj[to[i]]>ans)
            {
                ans=f[x]+f[to[i]]+bj[to[i]];
                p1=pos[x];p2=pos[to[i]];
            }
            if(f[to[i]]+bj[to[i]]>f[x])
            {
                f[x]=f[to[i]]+bj[to[i]];
                pos[x]=pos[to[i]];
            }
        }
}
int main()
{
    //freopen("tt.in","r",stdin);
    scanf("%d%d",&n,&K);
    for(int i=1,x,y;i<n;i++)
    {
        scanf("%d%d",&x,&y);
        add(x,y);add(y,x);
    }
    for(int i=1;i<=n;i++)bj[i]=1;
    dfs(1,0);fin+=ans;
    if(K==2)
    {   
        if(deep[p1]<deep[p2])swap(p1,p2);
        while(deep[p1]>deep[p2])
            bj[p1]=-1,p1=fa[p1];
        while(p1!=p2)
        {
            bj[p1]=-1;bj[p2]=-1;
            p1=fa[p1];p2=fa[p2];
        }
        ans=-n;
        dfs(1,0);fin+=ans;
    }
    printf("%d\n",2*(n-1)-fin+K);
    return 0;
}