思路：当 $k = 2$ 时，将直径上每条边边权修改为 $-1$ 后，再跑一遍直径。

但是我用求直径的两种方法得到的答案却不一样。

65pts code

#include <cstdio>
#include <algorithm>
#include <map>
#include <stack>

using namespace std;

const int N = 1e5 + 5;

int n, k, tot, head[N], dep[N], pos1, pos2, ans;
struct Edge {
	int to, next;
} edge[N << 1];
map<pair<int, int>, bool> ma;
stack<pair<int, int> > st;

inline void add(int u, int v) {
	edge[++tot].to = v;
	edge[tot].next = head[u], head[u] = tot;
}

inline void dfs(int x, int u, int &p) {
	if (dep[u] > dep[p])
		p = u;
	for (int i = head[u]; i; i = edge[i].next) {
		int v = edge[i].to;
		if (v == x)
			continue;
		if (!ma[make_pair(u, v)])
			dep[v] = dep[u] + 1;
		else
			dep[v] = dep[u] - 1;
		dfs(u, v, p);
	}
}

inline bool dfs_(int x, int u) {
	if (u == pos2) {
		while (!st.empty()) 
			ma[make_pair(st.top().first, st.top().second)] = ma[make_pair(st.top().second, st.top().first)] = true, st.pop();
		return true;
	}
	for (int i = head[u]; i; i = edge[i].next) {
		int v = edge[i].to;
		if (v == x)
			continue;
		st.push(make_pair(u, v));
		if (dfs_(u, v))
			return true;
		st.pop();
	}
	return false;
}

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(0, 1, pos1);
	dep[pos1] = 0;
	dfs(0, pos1, pos2);
	ans = dep[pos2];
	if (k == 2) {
		dfs_(0, pos1);
		dep[1] = pos1 = pos2 = 0;
		dfs(0, 1, pos1);
		dep[pos1] = 0;
		dfs(0, pos1, pos2);
		ans += dep[pos2];
	}
	printf("%d\n", 2 * n - ans - (k == 1));
	return 0;
}

AC code

#include <cstdio>
#include <algorithm>
#include <map>
#include <stack>

using namespace std;

const int N = 1e5 + 5;

int n, k, tot, head[N], dep[N], pos1, pos2, ans, res, dis[N];
struct Edge {
	int to, next;
} edge[N << 1];
map<pair<int, int>, bool> ma;
stack<pair<int, int> > st;

inline void add(int u, int v) {
	edge[++tot].to = v;
	edge[tot].next = head[u], head[u] = tot;
}

inline void dfs(int x, int u, int &p) {
	if (dep[u] > dep[p])
		p = u;
	for (int i = head[u]; i; i = edge[i].next) {
		int v = edge[i].to;
		if (v == x)
			continue;
		if (!ma[make_pair(u, v)])
			dep[v] = dep[u] + 1;
		else
			dep[v] = dep[u] - 1;
		dfs(u, v, p);
	}
}

inline bool dfs_(int x, int u) {
	if (u == pos2) {
		while (!st.empty()) 
			ma[make_pair(st.top().first, st.top().second)] = ma[make_pair(st.top().second, st.top().first)] = true, st.pop();
		return true;
	}
	for (int i = head[u]; i; i = edge[i].next) {
		int v = edge[i].to;
		if (v == x)
			continue;
		st.push(make_pair(u, v));
		if (dfs_(u, v))
			return true;
		st.pop();
	}
	return false;
}

inline void dfs1(int x, int u) {
	for (int i = head[u]; i; i = edge[i].next) {
		int v = edge[i].to;
		if (v == x)
			continue;
		dfs1(u, v);
		if (!ma[make_pair(u, v)]) {
			res = max(res, dis[u] + dis[v] + 1);
			dis[u] = max(dis[u], dis[v] + 1);
		} else {
			res = max(res, dis[u] + dis[v] - 1);
			dis[u] = max(dis[u], dis[v] - 1);
		}
	}
}

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(0, 1, pos1);
	dep[pos1] = 0;
	dfs(0, pos1, pos2);
	ans = dep[pos2];
	if (k == 2) {
		dfs_(0, pos1);
		dfs1(0, 1);
		ans += res;
	}
	printf("%d\n", 2 * n - ans - (k == 1));
	return 0;
}

一组 Hack 数据

10 2
6 10
5 3
6 4
7 3
8 3
6 1
1 9
2 1
6 3