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