Link
挺有意思的一道题,思路其实蛮清晰也比较容易理解,然而实现的时候写了半天
代码
const int maxn = 1e6 + 10;
const int maxm = 1e6 + 10;
struct Edge {
int to, dis, next;
}edge[maxm];
int n, m;
int head[maxn], dis[maxn], cnt = 1;
bool vis[maxn];
void add_edge(int u, int v, int w) {
cnt++;
edge[cnt].to = v;
edge[cnt].dis = w;
edge[cnt].next = head[u];
head[u] = cnt;
}
int k;
int fa[maxn];
int now;
void dfs(int cur) { //now表示目前的最大深度
vis[cur] = 1;
for(int i = head[cur]; i; i = edge[i].next) {
int v = edge[i].to;
if(vis[v]) continue;
dis[v] = dis[cur] + edge[i].dis;
if(dis[v] > dis[now]) {
now = v;
}
fa[v] = i;
dfs(v);
}
vis[cur] = 0;
}
int diam(int s = 1) {
now = s;
dfs(s);
memset(dis, 0, sizeof(dis));
memset(vis, 0, sizeof(vis));
memset(fa, 0, sizeof(fa));
dfs(now);
return dis[now];
}
void dp(int x) {
vis[x] = 1;
for(int i = head[x]; i; i = edge[i].next) {
int v = edge[i].to;
if(vis[v]) continue;
dp(v);
now = max(now, dis[x] + dis[v] + edge[i].dis);
dis[x] = max(dis[x], dis[v] + edge[i].dis);
}
}
/*
int dp(int i) {
if(vis[i]) return dis[i];
dis[i] = 0;
vis[i] = 1;
for(int j = head[i]; j; j = edge[j].next) {
int v = edge[j].to;
dis[i] = max(dis[i], dp(v) + edge[i].dis);
}
return dis[i];
}
*/
void solve() {
cin >> n >> k;
for(int i = 1; i < n; i++) {
int u, v;
cin >> u >> v;
add_edge(u, v, 1);
add_edge(v, u, 1);
}
int len = diam();
int ans = 2*(n-1)-len+1;
if(k == 2) {
while(fa[now]) {
edge[fa[now]].dis = edge[fa[now]^1].dis = -1;
now = edge[fa[now] ^ 1].to;
}
memset(dis, 0, sizeof(dis));
memset(vis, 0, sizeof(vis));
now = 0;
dp(1);
int maxx = 0;
for(int i = 1; i <= n; i++)
maxx = max(maxx, dis[i]);
// ans = ans - maxx + 1;
ans -= now - 1;
}
cout << ans << endl;
}