Codeforces 627D 题解

Codeforces 627D Preorder Test

思路

“前 \(k\) 个”“最小值最大”等字眼提示我们考虑二分最小权值 \(\operatorname{goal}\)。

为了方便表述，对树中任意一个结点 \(u\) 作如下几点定义：

• 将“以 \(u\) 为根的子树”简记为“子树 \(u\)”；

• 若 \(a_u \ge \operatorname{goal}\) 则称 \(u\) 为“黑点”，否则称 \(u\) 为“白点”；

• 若子树 \(u\) 内全是黑点，则称子树 \(u\) 是“满的”，否则称其是“不满的”。

那么我们的任务就是构造一个 DFS 序使得“遇到第一个白点前能遇到尽可能多的黑点”。DFS 时会一棵子树一棵子树地遍历，我们可以依此得出重要性质：如果一棵树的根已经确定（即 DFS 的起点确定），在 DFS 时我们必定会先遍历所有满的子树，然后选一个最优的不满子树进行遍历，直到碰到这个子树的第一个白点。

容易发现这个问题具有最优子结构性质——我们需要知道每个子树中“遇到第一个白点前能遇到多少黑点”。于是树形 DP 的状态呼之欲出——设 \(f(u)\) 表示：在子树 \(u\) 中，所有以 \(u\) 为起点的 DFS 序里遇到第一个白点前最多能遇到多少黑点。并且我们发现：子树 \(u\) 是满的当且仅当 \(f(u) = \operatorname{size}(u)\)，其中 \(\operatorname{size}(u)\) 表示子树 \(u\) 中结点的数量。

根据得出的重要性质，有转移方程：

\[f(u) = \begin{cases} \sum_{v \in \operatorname{son}(u) \ \wedge \ f(v) = \operatorname{size}(v)} f(v) + \max_{v \in \operatorname{son}(u) \ \wedge \ f(v) \ne \operatorname{size}(v)} \{ f(v) \} + 1 & a_u \ge \operatorname{goal} \\ 0 & a_u < \operatorname{goal} \end{cases} \]

如果这么做的话需要枚举整棵树的根结点（即换根），总时间复杂度为 \(\operatorname{O}(n^2 \log \max \{a \})\)，不可接受。我们希望优化掉换根的部分。于是我们设 \(g(u)\) 表示：在子树 \(u\) 中，所有经过 \(u\) 的 DFS 序里遇到第一个白点前最多能遇到多少黑点。为了使 \(g(u)\) 最大，DFS 一定是从 \(u\) 的某个不满子树开始，然后经过 \(u\)，再遍历 \(u\) 所有满的子树，最后再进入 \(u\) 的另一个不满子树。那么 \(g(u)\) 可以由 \(f(u)\) 直接推出来——它就是 \(f(u)\) 加上次优不满子树的 \(f\)。

但是还有一种情况：如果 \(u\) 没有不满子树或只有 1 个不满子树怎么办？前者是简单的，因为所有点都是黑点，都可以取。对于后者，我们可以把这个不满子树的贡献从 \(f\) 改为 \(g\)，这样这个不满子树就贡献了 DFS 的两端。

总时间复杂度为 \(\operatorname{O}(n \log \max \{a \})\)，空间复杂度为 \(\operatorname{O}(n)\)，可以接受。

实现

#include <iostream>
#include <vector>

using namespace std;

const int N = 2e5, A = 1e6;

int n, k, goal;
int a[N + 5], siz[N + 5], f[N + 5], g[N + 5];
vector<int> edge[N + 5];

void get_siz(int u, int fa) // 预处理 siz
{
    siz[u] = 1;
    for (auto v : edge[u])
        if (v != fa)
        {
            get_siz(v, u);
            siz[u] += siz[v];
        }
}

void calc(int u, int fa) // 树形 DP
{
    int max1 = 0, max2 = 0, maxg = 0; // 在所有不满子树中，最大的 f、次大的 f、最大的 g
    for (auto v : edge[u])
        if (v != fa)
            calc(v, u);
    if (a[u] >= goal) // 分类讨论，若 a[u] < goal 则显然 f[u] = g[u] = 0，在初始化中已实现
    {
        f[u] = 1; // u 自己是黑点
        for (auto v : edge[u])
            if (v != fa)
                if (f[v] == siz[v]) // 子树 v 是满的
                    f[u] += f[v];
                else // 子树 v 是不满的，维护 max1, max2, maxg
                {
                    if (f[v] >= max1)
                        max2 = max1, max1 = f[v];
                    else
                        max2 = max(max2, f[v]);
                    maxg = max(maxg, g[v]);
                }
        // 求解 g[u]（当前 f[u] 未加上不满子树的贡献）
        if (max1 == 0) // 子树全满
            g[u] = f[u];
        else if (max2 == 0) // 只有 1 个不满子树
            g[u] = f[u] + maxg;
        else // 有不止 1 个不满子树
            g[u] = f[u] + max1 + max2;
        f[u] += max1; // 为 f[u] 加上不满子树的贡献
    }
}

bool check(int val) // 二分中的 check()
{
    int res = 0;
    goal = val;
    // 树形 DP 初始化
    fill(f + 1, f + n + 1, 0);
    fill(g + 1, g + n + 1, 0);
    calc(1, 0);
    for (int i = 1; i <= n; i++)
        res = max(res, g[i]);
    return res >= k;
}

int search() // 二分函数
{
    int l = 1, r = A, res = -1;
    while (l <= r)
    {
        int mid = (l + r) / 2;
        if (check(mid))
            l = mid + 1, res = mid;
        else
            r = mid - 1;
    }
    return res;
}

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);

    int maxa = 0;
    cin >> n >> k;
    for (int i = 1; i <= n; i++)
        cin >> a[i];
    for (int i = 1, u, v; i < n; i++)
    {
        cin >> u >> v;
        edge[u].push_back(v);
        edge[v].push_back(u);
    }
    get_siz(1, 0);
    cout << search() << "\n";

    return 0;
}