P3629 [APIO2010]巡逻(树的直径 树形dp)

当k=1时答案就是树的直径头尾相连
当k=2时答案就是把树的直径上的点标记起来并将相邻的两个直径上的点的距离标记为-1，再用树形dp找树的直径

#include<bits/stdc++.h>
//#define int long long
#define memarray(array, value) memset(array, value, sizeof(array))
using namespace std;
const double EPS=1e-5;
const double PI=acos(-1);
const long long mod=1e9+7;
const int INF=0x3f3f3f3f;
inline int sgn(double a){ return a < -EPS ? -1 : a > EPS; }
inline int cmp(double a, double b){ return sgn(a-b); }
const int MAXN=5e5+10;
int n,m;
vector<int>vec[MAXN];
int dis[MAXN];
int head,tail;
int fa[MAXN],huan[MAXN];
void dfs(int now,int pre){
    for(auto to:vec[now]){
        if(to==pre)continue;
        dis[to]=dis[now]+1;
        fa[to]=now;
        dfs(to,now);
    }
}
int Dafter;
void dfs2(int now,int pre){
    for(auto to:vec[now]){
        if(to==pre)continue;
        dfs2(to,now);
        if(huan[to]&&huan[now]){
            Dafter=max(Dafter,dis[now]+dis[to]-1);
            dis[now]=max(dis[now],dis[to]-1);
        }else{
            Dafter=max(Dafter,dis[now]+dis[to]+1);
            dis[now]=max(dis[now],dis[to]+1);
        }
    }
}
void solve(){
    dfs(1,0);//找直径的头
    head=1;
    for(int i=2;i<=n;i++)
        if(dis[head]<dis[i])
            head=i;
    dis[head]=0;fa[head]=head;
    dfs(head,0);//找直径的尾
    tail=1;
    for(int i=2;i<=n;i++)
        if(dis[tail]<dis[i])
            tail=i;
    if(m==1){
        printf("%d\n",(n-1)*2-dis[tail]+1);
        return;
    }
    int d1=(n-1)*2-dis[tail]+1;
    int now=tail;//将直径上的点标记起来，直径上的点到相邻的直径上的点的距离为-1
    while(fa[now]!=now){
        huan[now]=true;
        now=fa[now];
    }
    huan[now]=true;
    memarray(dis,0);
    dfs2(1,0);//树形dp找带有负边权的树的最大直径
    printf("%d\n",d1-Dafter+1);
}
void init(){
    scanf("%d%d",&n,&m);
    for(int i=1,u,v;i<n;i++){
        scanf("%d%d",&u,&v);
        vec[u].push_back(v);
        vec[v].push_back(u);
    }
}
signed main(){
    int T;
//    scanf("%d",&T);
    T=1;
    while(T--){
//    for(int i=1;i<=T;i++){
        init();
        solve();
    }
    return 0;
}