Codeforces Round #495 (Div. 2) E. Sonya and Ice Cream 树形dp

E. Sonya and Ice Cream

题意：给一棵树，你可以选择一条长度为 $k$的链，使得所有点与链的最大距离最小，求出最小距离。

解法：我们先用换根 $dp$（直接树形 $dp$也行，就是想秀一波换根 $dp$）求出距离所有点最大距离最小的根 $rt$，然后以 $rt$为根树形 $dp$求出每个点到最远的叶子节点距离 $d[u]$，定义 $dp[u]$为 $d[u] + dist(u, father[u])$，我们找到 $rt$所有儿子中 $dp[son]$最大的那两个，然后加入优先队列，每次出队都把该节点 $dp[son]$最大的儿子加入优先队列，直到出 $k-1$次队列，所有出队的节点就是我们要选择的链，然后统计答案即可。

#include<bits/stdc++.h>
#define pi pair<int, int>
#define mk make_pair
using namespace std;
const int maxn = 1e5 + 10;
vector<pi> G[maxn];
vector<int> suf[maxn];
int d[maxn], f[maxn], mn = 2e9, rt;
void dfs(int u, int fa) {
    d[u] = 0;
    f[u] = fa;
    for (auto tmp : G[u])
        if (tmp.first != fa) {
            int v = tmp.first;
            dfs(v, u);
            d[u] = max(d[u], d[v] + tmp.second);
        }
}
void dfs2(int u, int fa) {
    d[u] = 0;
    for (auto tmp : G[u])
        d[u] = max(d[u], d[tmp.first] + tmp.second);
    if (d[u] < mn)
        mn = d[u], rt = u;
    suf[u].push_back(0);
    for (int i = G[u].size() - 1; i > 0; i--) {
        pi tmp = G[u][i];
        int val = max(suf[u].back(), d[tmp.first] + tmp.second);
        suf[u].push_back(val);
    }
    int mx = 0;
    for (auto tmp : G[u]) {
        d[u] = max(mx, suf[u].back());
        mx = max(mx, d[tmp.first] + tmp.second);
        suf[u].pop_back();
        if (tmp.first != fa)
            dfs2(tmp.first, u);
    }
}
priority_queue<pi> q;
int main() {
    int n, k, u, v, w;
    scanf("%d%d", &n, &k);
    for (int i = 1; i < n; i++) {
        scanf("%d%d%d", &u, &v, &w);
        G[u].push_back(mk(v, w));
        G[v].push_back(mk(u, w));
    }
    dfs(1, 0);
    dfs2(1, 0);
    dfs(rt, 0);
    if (k == 1)
        return printf("%d\n", d[rt]), 0;
    pi Mx = mk(0, 0), Mx2 = mk(0, 0), Mx3 = mk(0, 0);
    for (auto tmp : G[rt]) {
        int val = d[tmp.first] + tmp.second;
        if (val > Mx.first)
            Mx3 = Mx2, Mx2 = Mx, Mx = mk(val, tmp.first);
        else if (val > Mx2.first)
            Mx3 = Mx2, Mx2 = mk(val, tmp.first);
        else if (val > Mx3.first)
            Mx3 = mk(val, tmp.first);
    }
    q.push(Mx);
    q.push(Mx2);
    int ans = Mx3.first;
    for (int i = 2; i <= k; i++) {
        u = q.top().second;
        q.pop();
        pi mx = mk(0, 0), mx2 = mk(0, 0);
        for (auto tmp : G[u])
        if (tmp.first != f[u]) {
            int val = d[tmp.first] + tmp.second;
            if (val > mx.first)
                mx2 = mx, mx = mk(val, tmp.first);
            else if (val > mx2.first)
                mx2 = mk(val, tmp.first);
        }
        ans = max(ans, mx2.first);
        q.push(mx);
    }
    printf("%d\n", max(ans, q.top().first));
}