【树的直径】洛谷_3629 [APIO2010]巡逻

版权声明：本文为博主原创文章，未经博主允许不得转载。 https://blog.csdn.net/SSL_hzb/article/details/85011322

题意

有一颗树，要在其中加入 $K(k\leq 2)$条边，使得本来遍历这颗树要经过的边数最少，同时加入的边一定要正好走过 $1$次。

思路

当 $K=1$时，显然是在树的直径的两个点之间连一条边，因为这样可以少走一次直径，而直径又是最长的。
当 $K=2$时，新建的边如果与之前的边没有环重叠的话也是和 $K=1$的方法一样。如果有重叠的话，那么要正好走过一次新建的边，我们重复的边就要都多走一次。
综上所述，可以这样做:
$1)$先求一次树的直径，记为 $D_1$，然后把直径上的边取反。
$2)$再求一次直径，记为 $D_2$。
答案为 $2(N-1)-(D_1-1)-(D_2-1)$，如果我们取到重复的边，因为它是负的，根据初中数学，可以知道它会加回来。

代码

#include<cstdio>
#include<cstring>
#include<algorithm>

struct node{
    int to, next, v;
}e[200001];
int N, K, tot = 1, l, r, d;
int head[100001], pre[100001], v[100001], f[100001];

void add(int x, int y) {
    e[++tot].to = y;
    e[tot].next = head[x];
    e[tot].v = 1;
    head[x] = tot;
}

void dfs1(int p, int fa, int w) {
    if (w >= d) {
        d = w;
        l = p;
    }
    for (int i = head[p]; i; i = e[i].next) {
        if (e[i].to != fa)
            dfs1(e[i].to, p, w + e[i].v);
    }
}

void dfs2(int p, int fa, int w) {
    if (w >= d) {
        d = w;
        r = p;
    }
    for (int i = head[p]; i; i = e[i].next) {
        if (e[i].to != fa) {
            pre[e[i].to] = p;
            dfs2(e[i].to, p, w + e[i].v);
        }
    }
}

void dfsD() {//dfs求树的直径
    d = 0;
    dfs1(1, 0, 0);
    d = 0;
    dfs2(l, 0, 0);
}

void change(int p) {//暴力回溯改边
    if (p == l) return;
    for (int i = head[p]; i; i = e[i].next) {
        if (e[i].to == pre[p]) {
            e[i].v = -1;
            e[i ^ 1].v = -1;
            change(e[i].to);
        }
    }
}

void dp(int x) {
    v[x] = 1;
    for (int i = head[x]; i; i = e[i].next) {
        if (v[e[i].to]) continue;
        dp(e[i].to);
        d = std::max(d, f[x] + f[e[i].to] + e[i].v);
        f[x] = std::max(f[x], f[e[i].to] + e[i].v);
    }
}

int main() {
    scanf("%d %d", &N, &K);
    for (int i = 1, a, b; i < N; i++) {
        scanf("%d %d", &a, &b);
        add(a, b);
        add(b, a);
    }
    dfsD();
    if (K == 1) {
        printf("%d", 2 * N - 1 - d);
        return 0;
    } else {
        int s = 2 * N - d;
        change(r);
        d = 0;
        dp(1);//边权有负所以用树形dp
        printf("%d", s - d);
        return 0;
    }
}