D50 树的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 u,int v){
  to[++idx]=v,w[idx]=1,ne[idx]=h[u],h[u]=idx;
}
int d[N],pre[N],col[N];
int n,k,p,d1,d2;

void dfs(int u,int fa){
  if(d[u]>d[p]) p=u; //记录直径端点
  pre[u]=fa;         //记录直径路径
  for(int i=h[u];i;i=ne[i]){
    int v=to[i];
    if(v!=fa){
      d[v]=d[u]+w[i]; //记录v到根的距离
      dfs(v,u);
    }
  }
}
void dfs2(int u,int fa){
  for(int i=h[u];i;i=ne[i]){
    int v=to[i];
    if(v!=fa){
      if(col[u] && col[v]) w[i]=-1; //直径边取反
      dfs2(v,u);
      d2=max(d2,d[u]+d[v]+w[i]); //记录经过u的直径
      d[u]=max(d[u],d[v]+w[i]);  //记录u的最长链长度
    }
  }
}
int main(){
  scanf("%d%d",&n,&k);
  for(int i=1,u,v;i<n;++i){
    scanf("%d%d",&u,&v);
    add(u,v),add(v,u);
  }
  dfs(1,0); d[p]=0;
  dfs(p,0); d1=d[p]; //两次DFS求直径d1，记录直径路径pre
  
  if(k==1){
    cout<<2*(n-1)-d1+1;
    return 0;
  }
  for(int i=p;i;i=pre[i])col[i]=1; //直径染色
  memset(d,0,sizeof d);
  dfs2(1,0); //树形DP求直径d2
  cout<<2*(n-1)-(d1+d2)+2;
}

/* by: hansang（scy修改） 不建新边时，要走 2*(n-1)次，每条道路经过两次 初始化ans=2*(n-1) (1) 当 K=1 随机建一条边，必定形成一个环 环里的每条旧边只用走一次，环越大答案越小 找出树的直径，把直径的头尾连边 设不建新边时树的直径长度为 dis，则有当 K=1，答案为 ans-dis+1 (因为新边也走了一次，之前的 ans没算，要加上1)

(2) 当 K=2 假设像之前一样求直径，如果两条直径有重叠 则重叠部分里的边被减了两次，相当于重叠部分的点没走过。 实际上重叠的边要走两次。所以求第二条路径前，把第一条直径的边权为-1 设第二次求树的直径为 dis，答案为 ans-dis+1

求直径的方法： 第一遍求直径时，因为要记录边，使用2遍 dfs1别找到直径的两端（分选任意一点为根， dfs1找到最深的点L，然后让L为根，dfs1找到最深的点R，L到R的距离为直径之一）， 第二遍求直径时，边权有负数不能使用 dfs1，则考虑使用 dfs2 */

#include <bits/stdc++.h>
using namespace std;
const int N=1e5+10;
struct edge{int x, y, c, pre;} a[2*N]; int alen, last[N];
void ins(int x, y, int c) {alen++; a[alen]=edge{x, y, c, last[x]}; last[x]=alen;}
bool v[N]; int dis, fa, f[N], b[N], d[N];
void dfs1(int x)
{
	v[x]=1; //记录，防止重复经过一个点 
	for(int k=last[x]; k; k=a[k].pre)
	{
		int y=a[k].y; if(v[y]==1) continue;
		d[y]=d[x]+a[k].c; f[y]=x; b[y]=k; //赋值，记录点和边 
		if(dis < d[y]) dis=d[y], fa=y; //更新直径长度 
		dfs1(y);
	}
}
void dfs2(int x)
{
	v[x]=1;
	for(int k=last[x]; k; k=a[k].pre) 
	{
		int y=a[k].y; if(v[y]==1) continue;
		dfs2(y); //注意下面这两句的顺序 
		dis=max(dis, d[x]+d[y]+a[k].c); //以 x为顶点的直径 
		d[x]=max(d[x], d[y]+a[k].c); //用 y去更新 x往下最长的链 
	}
	d[x]=max(d[x], 0); //当 x没有孩子节点 或 d[x]为负数 
	dis=max(dis, d[x]);
}
int main()
{
	int n, K; scanf("%d%d", &n, &K);
	int ans=(n-1)<<1; // ans=2*(n-1) 
	alen=1; memset(last, 0, sizeof(last));
	for(int i=1; i<n; i++)
	{
		int x, y; scanf("%d%d", &x, &y); //这里原代码可能有笔误，应为&x, &y
		ins(x, y, 1); ins(y, x, 1); //边权为 1 
	}
	int L, R; 
	memset(v, 0, sizeof(v));
	d[1]=0;f[1]=0; dis=-N; dfs1(1); L=fa; //任选一点出发，记录最长的链的另一个端点L 
	memset(v, 0, sizeof(v)); 
	d[L]=0;f[L]=0; dis=-N; dfs1(L); R=fa; //从L出发，记录最长的链的另一个端点R，这样就得到直径的两个端点了 
	
	ans=ans-dis+1; //计算 ans 
	
	if(K==2)
	{
		memset(v, 0, sizeof(v));
		memset(d, -63, sizeof(d)); //有负权边，d赋值负无穷 
		for(int i=R; i; i=f[i]) //从 R开始，遍历直径上的点 
		{
			a[b[i]].c=-1; //把边取反 
			a[b[i]^1].c=-1; //另一条反向边 
		}
		dis=0; dfs2(1); ans=ans-dis+1; //计算 ans
	}
	printf("%d\n", ans);
	return 0;
}