P3629巡逻 树的直径*

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;
}