Author: lllyouo
Date: 20250305
tag: 树的直径
link: https://www.acwing.com/problem/content/352/

问题描述

link

分析

略

参考代码

#include <bits/stdc++.h>
using namespace std;

const int N = 1e5 + 10, M = 2 * N;
// idx 从 2 开始，成对变换
int h[N], e[M], w[M], ne[M], idx = 2;
int n, k;
int d[N], prv[N], L;
queue<int> q;

void add(int a, int b, int c) {
	e[idx] = b;
	w[idx] = c;
	ne[idx] = h[a];
	h[a] = idx++;
}

int bfs(int s) {
	memset(d, -1, sizeof d);
	q.push(s);
	d[s] = 0;
	
	while (q.size()) {
		int x = q.front();
		q.pop();
		for (int i = h[x]; i != -1; i = ne[i]) {
			int y = e[i];
			if (d[y] == -1) {
				d[y] = d[x] + 1;
				prv[y] = i;
				q.push(y);
			}
		}
	}
	int p = s;
	for (int i = 1; i <= n; i++) {
		if (d[i] > d[p]) p = i;
	}
	return p;
}

void update(int q, int p) {
	while (q != p) {
		w[prv[q]] = -1;
		w[prv[q] ^ 1] = -1;
		q = e[prv[q] ^ 1];
	}
}

void dp(int x, int fx) {
	for (int i = h[x]; i != -1; i = ne[i]) {
		int y = e[i];
		if (y == fx) continue;
		dp(y, x);
		L = max(L, d[y] + d[x] + w[i]);
		d[x] = max(d[x], d[y] + w[i]);
	}
}

int main() {
    memset(h, -1, sizeof h);
	cin >> n >> k;
	for (int i = 1; i < n; i++) {
		int a, b; cin >> a >> b;
		add(a, b, 1);
		add(b, a, 1);
	}
	
	int p = bfs(1);
	int q = bfs(p);
	int ans = 2 * (n - 1) - d[q] + 1;
	
	if (k == 1) cout << ans << endl;
	else {
		update(q, p);
		memset(d, 0, sizeof d);
		dp(1, 0);
		cout << ans - L + 1 << endl;
	}
	
	return 0;
}