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,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); d[p]=0;
  dfs(p,0); d1=d[p]; //兩次DFS
  
  for(int i=p;i;i=fa[i])col[i]=1; //直徑染色
  change(p,0); //修改直徑的邊權
  memset(d,0,sizeof d);
  dfs2(p,0);   //樹形DP
  
  if(k==1) printf("%d\n",2*n-d1-1);
  else printf("%d\n",2*n-d1-d2);
}