bzoj1912 树形dp求直径(新写法)，求直径的两端点

通过回溯法可以求出直径的两个端点，同时注意有负权边的树求直径不可以用两次dfs来求，而必须用dp做

/*
分情况讨论问题
一条边也不加的情况，显然每条边要扫描两次，
    该情况的答案是2(n-1)
只加一条边的情况，找到直径，将其变成一个环，在这个环上的所有边只要扫描一次，剩下的边就要扫描两次
    设直径为L，该情况下的答案是 2(n-1-L)+L+1=2n-L-1=2(n-1)-(L-1)
加两条边的情况，在加入第一条边出现环的情况下，再加入一条边形成的环会和原来的环有重合
重合的部分任然要扫描两次，所以新加入的边要使形成的环和原来的环非重合的边和重合的边的差最大
那么如何快速求新加入的边？只要将第一条边加入后形成的环的边权值变成-1，然后再求一次直径即可
    设第一次的直径L1，第二次的直径L2，那么该情况下的答案是2(n-1)-(L1-1)-(L2-1) 
*/
#include<bits/stdc++.h>
using namespace std;
#define maxn 100005
struct Edge{int to,nxt,w;}edge[maxn<<1];
int head[maxn],tot,n,k,ans;
void init(){
    memset(head,-1,sizeof head);
    tot=0;
}
void addedge(int u,int v){
    edge[tot].w=1;edge[tot].to=v;edge[tot].nxt=head[u];head[u]=tot++;
}

int pre1[maxn],pre2[maxn],dp[maxn],mx;//记录x的最长和次长子链的边的下标 
int dfs(int x,int fa)
{
    int mx1=0,mx2=0;
    for(int i=head[x];i!=-1;i=edge[i].nxt)
        if(edge[i].to!=fa)
        {
            int v=edge[i].w+dfs(edge[i].to,x);
            if(v>mx1)mx2=mx1,mx1=v,pre2[x]=pre1[x],pre1[x]=i;
            else if(v>mx2)mx2=v,pre2[x]=i;
        }
    if(mx1+mx2>ans)ans=mx1+mx2,mx=x;
    return mx1;
}

int main(){
    init();
    memset(pre1,-1,sizeof pre1);
    memset(pre2,-1,sizeof pre2);
    cin>>n>>k;
    int tot=2*n-2;
    for(int i=1;i<n;i++){
        int u,v;cin>>u>>v;
        addedge(u,v);addedge(v,u);
    }    
    dfs(1,0);
    tot=tot-ans+1;
    if(k==2)
    {
        ans=0;
        for(int i=pre1[mx];i!=-1;i=pre1[edge[i].to])edge[i].w=edge[i^1].w=-1;
        for(int i=pre2[mx];i!=-1;i=pre1[edge[i].to])edge[i].w=edge[i^1].w=-1;
        dfs(1,0);tot=tot-ans+1;
    }
    cout<<tot<<endl;
}