2018.11.06 bzoj1912: [Apio2010]patrol 巡逻（树形dp）

传送门
一道挺妙的题啊。
对于$K==1$的直接求树的直径。
对于$K==2$的先求一次直径，然后考虑到如果两条边加进去形成的两个环重叠就会有负的贡献。
因此把之前那条直径上的边权改成-1再求一次直径就可以了。
代码：

#include<bits/stdc++.h>
using namespace std;
inline int read(){
    int ans=0;
    char ch=getchar();
    while(!isdigit(ch))ch=getchar();
    while(isdigit(ch))ans=(ans<<3)+(ans<<1)+(ch^48),ch=getchar();
    return ans;
}
const int N=1e5+5;
int ans,n,K,bg=0,ed=0,first[N],cnt=0,dis[N],fa[N],dismax;
bool vis[N];
struct Edge{int u,v;}tt[N];
struct edge{int v,next,w;}e[N<<1];
inline void add(int u,int v,int w){e[++cnt].v=v,e[cnt].next=first[u],e[cnt].w=w,first[u]=cnt;}
inline void dfs1(int p,int pre){
    for(int i=first[p];i;i=e[i].next){
        int v=e[i].v;
        if(v==pre)continue;
        dis[v]=dis[p]+e[i].w,dfs1(v,p);
    }
    if(dis[p]>dis[bg])bg=p;
}
inline void dfs2(int p){
    for(int i=first[p];i;i=e[i].next){
        int v=e[i].v;
        if(v==fa[p])continue;
        fa[v]=p,dis[v]=dis[p]+e[i].w,dfs2(v);
    }
    if(dis[p]>dis[ed])ed=p;
}
inline void dfs(int p,int pre){
    for(int i=first[p];i;i=e[i].next){
        int v=e[i].v;
        if(v==pre)continue;
        dfs(v,p),dismax=max(dismax,dis[p]+dis[v]+e[i].w),dis[p]=max(dis[p],dis[v]+e[i].w);
    }
}
int main(){
    n=read(),K=read();
    for(int i=1,u,v;i<n;++i)u=read(),v=read(),add(u,v,1),add(v,u,1),tt[i]=(Edge){u,v};
    dfs1(1,0),dis[bg]=0,dfs2(bg);
    if(K==1)return printf("%d",n*2-1-dis[ed]),0;
    ans=n*2-dis[ed],memset(first,0,sizeof(first)),cnt=0;
    for(int i=ed;i;i=fa[i])vis[i]=1;
    for(int i=1,u,v;i<n;++i)u=tt[i].u,v=tt[i].v,add(u,v,vis[u]&&vis[v]?-1:1),add(v,u,vis[u]&&vis[v]?-1:1);
    fill(dis+1,dis+n+1,0),dfs(1,0);
    cout<<ans-dismax;
    return 0;
}