D50 树的直径 两次DFS+树形DP P3629 [APIO2010] 巡逻

D50 树的直径 P3629 [APIO2010] 巡逻_哔哩哔哩_bilibili 

P3629 [APIO2010] 巡逻 - 洛谷

// 两次DFS+树形DP O(n)
#include<bits/stdc++.h>
using namespace std;

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

void dfs(int x,int f){
  if(d[x]>d[p]) p=x; //更新端点  
  pre[x]=f; //记录路径
  for(int i=h[x];i;i=ne[i]){
    int y=to[i];
    if(y==f)continue;
    d[y]=d[x]+w[i];
    dfs(y,x);
  }
}
void dp(int x,int f){
  for(int i=h[x];i;i=ne[i]){
    int y=to[i];
    if(y==f)continue;
    if(col[x]&&col[y]) w[i]=-1;
    dp(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];
  if(k==1){
    printf("%d\n",2*n-d1-1);
    return 0;
  }
  
  for(int i=p;i;i=pre[i])col[i]=1;
  memset(d,0,sizeof d);
  dp(1,0); //树形dp
  printf("%d\n",2*n-d1-d2);
}