CF627D Preorder Test 题解

题意：

为了她的计算机科学课程，cxy 用小棍和小球构建了一个包含 \(n\) 个节点的树模型。cxy 花了 \(a_{i}\) 分钟构建树中的第 \(i\) 个小球。

cxy 的老师会评估她的模型，并根据她投入的努力给她打分。然而，老师没有足够的时间来检查整个树以确定这些。cxy 知道老师将会检查树的前 \(k\) 个节点（按 dfs 序先序遍历树）。然后他会根据这 \(k\) 个节点中发现的最小 \(a_{i}\) 给 cxy 打分。

虽然 cxy 没有足够的时间重建她的模型，但她可以选择老师开始检查的根节点。此外，她可以以任何顺序重新排列每个节点的邻居列表。请帮助 cxy 在这次作业中获得最佳分数。

输入格式：

输入的第一行包含两个正整数，\(n\) 和 \(k~(2 \leq n \leq 2\times 10^5,~1 \leq k \leq n)\) ，表示 cxy 的树中的小球数量和老师将检查的小球数量。

第二行包含 \(n\) 个整数，\(a_{i}~(1 \leq a_{i} \leq 10^6)\)，cxy 用来构建第 \(i\) 个小球的时间。

接下来的 \(n-1\) 行中每行包含两个整数 \(u_{i}\) 和 \(v_{i}~(1 \leq u_{i}, v_{i} \leq n,\,u_{i} \ne v_{i})\)，表示 cxy 的树中小球 \(u_{i}\) 和 \(v_{i}\) 之间的连接。

输出格式：

输出一个整数，即 cxy 通过选择正确的树根并重新排列邻居列表所能获得的最高分数。

这里讲一种不需要换根的树形dp。

考虑二分最小值的最大值 \(\rm mid\) ，问题转化为是否存在长大于 \(k\) 的 dfs 序，

且 dfs 序中每个点权值 \(a_i\) 均不小于 \(\rm mid\) ，这里 dfs 路径的起点和终点随意。

考虑设 \(f(u)\) 表示 \(u\) 所在子树中以 \(u\) 为起点的满足 \(a_i \ge \mathrm{mid}\) 的最长 dfs 路径的长。

显然 \(a_u < \mathrm{mid}\) 时 \(f(u)=0\) 。考虑转移。对于 \(u\) 的任意儿子 \(v\) ，

若 \(f(v) = \operatorname{size} v\) ，那么可以直接把 \(f(v)\) 加到 \(f(u)\) 上，也就是往 \(v\) 转悠一圈再回到 \(u\) 。
若 \(f(v)<\operatorname{size} v\) ，我们先统计出这种情况的最大值和次大值（如果没有则为 \(0\)）。

在计算出所有 \(f(v) = \operatorname{size} v\) 的贡献后，显然以 \(u\) 开头的 dfs 路径还可以走到某一棵子树里。

于是我们选择 \(f(v)<\operatorname{size} v\) 中 \(f(v)\) 最大的那棵树，走到它里面就好了，那么有转移方程

\[ f(u) = \sum_{v \in A} f(v) + \max_{v \in B}\{f(v)\} \] 其中集合 \(A,B\) 分别表示 \(f\) 等于 \(\operatorname{size}\) 的 \(v\) 的集合以及 \(f\) 小于 \(\operatorname{size}\) 的 \(v\) 的集合。

不过显然在 \(u\) 子树内的最长 dfs 序不一定以 \(u\) 为起点，因此我们记录一个变量 \(\rm ans\) 用于二分的判断。

由于我们统计了次大值，所以实际上这条路径就是从次大值的树里走到 \(u\) ，再从 \(u\) 走最优方案的情况。

那么我们每次用 \(f(u)\) 加上次大值来更新 \(\rm ans\) ，最后判断 \(\rm ans\) 是否大于等于 \(k\) 即可……吗？

实际上这样的 dp 方式在一些情况下是错误的，原因在于我们没有考虑到 \(fa(u)\) 可以转悠回来的情况

换句话说，如果把 \(u\) 提作根重新 dp 一遍，原来的 \(fa(u)\) 满足 \(f_{\rm new}(fa(u)) = \operatorname{size}_{\rm new} fa(u)\) 。

怎么办呢？难道只能换根 dp 了吗？~~不过既然有这篇题解不就说明有办法么。~~

注意到，如果我们钦定 \(a_i\) 最小的那个节点为根，那么 \(f(\mathrm{root})\) 的取值只可能是 \(0\) 或者 \(n\) 。

显然 \(f(\mathrm{root})=0\) 的时候，任意的非根节点它的 \(fa\) 都不可能转回来，所以没有问题。

而当 \(f(\mathrm{root})=n\) 时，任何节点都有 \(f(u) = \operatorname{size} u\) ，然后 \(f(\mathrm{root})\) 也会被统计，所以也没有问题。

于是我们就不需要换根啦，找到 \(a_i\) 最小的点提作根直接 dp 即可。

时间复杂度 \(\mathcal{O}(n \log V)\)

代码：

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF 0x3f3f3f3f3f3f3f3f
void up(int &x, int y) { x < y ? x = y : 0; }
void down(int &x, int y) { x > y ? x = y : 0; }
#define rep(i, a, b) for(int i = (a), i##END = (b); i <= i##END; i++)
#define Rep(i, a, b) for(int i = (a), i##END = (b); i >= i##END; i--)
#define pb push_back
#define N ((int)(2e5 + 15))

vector<int> vec[N];
int rt, ans, mid, a[N], f[N], sz[N];
void dfs(int u, int fa)
{
    f[u] = sz[u] = 1; int mx = 0, smx = 0;
    for(auto v : vec[u]) if(v != fa)
    {
        dfs(v, u); sz[u] += sz[v];
        if(f[v] == sz[v]) f[u] += f[v];
        else
        {
            if(f[v] > mx) { smx = mx, mx = f[v]; }
            else if(f[v] > smx) smx = f[v];
        }
    }
    f[u] = a[u] < mid ? 0 : f[u] + mx; up(ans, f[u] + smx);
}
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);
    // freopen("check.in","r",stdin);
    // freopen("check.out","w",stdout);
    int n, k; cin >> n >> k; rep(i, 1, n) cin >> a[i];
    for(int i = 1, u, v; i < n; i++) { cin >> u >> v; vec[u].pb(v); vec[v].pb(u); }
    rt = min_element(a + 1, a + 1 + n) - a;
    int l = *min_element(a + 1, a + 1 + n), r = *max_element(a + 1, a + 1 + n);
    while(l < r)
    {
        mid = (l + r + 1) >> 1; ans = 0;
        if(dfs(rt, 0), ans >= k) l = mid; else r = mid - 1;
    }
    cout << l << '\n';
    return 0;
}