D50 树的直径 P3629 [APIO2010] 巡逻

视频链接：

P3629 [APIO2010] 巡逻 - 洛谷 | 计算机科学教育新生态 (luogu.com.cn)

// 两次DFS+树形DP O(n)
#include <iostream>
#include <cstring>
#include <algorithm>
using namespace std;

const int N=100005;
int idx=1,head[N],to[N<<1],ne[N<<1],w[N<<1];
void add(int x,int y){
  to[++idx]=y,w[idx]=1,ne[idx]=head[x],head[x]=idx;
}
int d[N],fa[N],col[N];
int n,k,p,st,ed,d1,d2;

void dfs(int x,int f){
  fa[x]=f; //记录路径
  if(d[x]>d[p]) p=x; //更新端点
  for(int i=head[x];i;i=ne[i]){
    int y=to[i];
    if(y==f)continue;
    d[y]=d[x]+w[i];
    dfs(y,x);
  }
}
void change(int x,int f){
  for(int i=head[x];i;i=ne[i]){
    int y=to[i];
    if(y==f)continue;
    if(col[y])w[i]=-1,w[i^1]=-1;
    change(y,x);
  }
}
void dfs2(int x,int f){
  for(int i=head[x];i;i=ne[i]){
    int y=to[i];
    if(y==f)continue;
    dfs2(y,x);
    d2=max(d2,d[y]+d[x]+w[i]);
    d[x]=max(d[x],d[y]+w[i]);
  }
}
int main(){
  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);
  }
  dfs(1,0); st=p; d[p]=0;
  dfs(p,0); ed=p; d1=d[p]; //两次DFS
  
  for(int i=ed;i;i=fa[i])col[i]=1; //直径染色
  change(st,0); //修改直径的边权
  memset(d,0,sizeof d);
  dfs2(st,0);   //树形DP
  
  if(k==1) printf("%d\n",2*n-d1-1);
  else printf("%d\n",2*n-d1-d2);
}