この問題は好き。でも時間内には無理。

E. Preorder Test

問題

木が与えられる。頂点には重みが振ってある。 この木と同型な根付き木からdfsの順に$k$個頂点を取ってきたとき、その重みの最小値の最大値を答えよ。

解法

二分探索。$O(n \log n)$。

重みが$m$未満の頂点の使用を禁止して$k$個使えるかを判定する。 禁止頂点によって分けられた各連結成分ごとに見る。 各頂点ごとに、その頂点を始点とし、その部分木内で最大何個使えるかを再帰的に求める。 禁止頂点を含まない部分木はただで使えるが、禁止頂点に隣接した頂点をその連結成分内の根としておけば、親方向の部分木がただで使えることがなくなり楽である。

実装

#include <iostream>
#include <vector>
#include <algorithm>
#include <functional>
#define repeat(i,n) for (int i = 0; (i) < (n); ++(i))
template <class T> bool setmax(T & l, T const & r) { if (not (l < r)) return false; l = r; return true; }
using namespace std;
int main() {
    int n, k; cin >> n >> k;
    vector<int> a(n); repeat (i,n) cin >> a[i];
    vector<vector<int> > g(n);
    repeat (i,n-1) {
        int u, v; cin >> u >> v; -- u; -- v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    int low  = *min_element(a.begin(), a.end());
    int high = *max_element(a.begin(), a.end()) + 1; // [low, high)
    while (low + 1 < high) {
        int mid = (low + high) / 2;
        auto forbidden = [&](int i) { return a[i] < mid; };
        bool found = false;
        vector<bool> used(n);
        vector<int> dp(n);
        vector<bool> has_forbidden(n);
        repeat (i,n) if (not forbidden(i) and not used[i]) {
            int root = -1;
            function<void (int)> collect = [&](int i) {
                used[i] = true;
                for (int j : g[i]) if (not used[j]) {
                    if (not forbidden(j)) collect(j);
                    if (forbidden(j)) root = i;
                }
            };
            collect(i);
            if (root == -1) {
                found = true;
                break;
            }
            function<void (int, int)> dfs = [&](int i, int p) {
                if (found) return;
                int max_forbidden = 0;
                int snd_forbidden = 0;
                int sum_free = 0;
                for (int j : g[i]) if (j != p) {
                    if (forbidden(j)) {
                        has_forbidden[i] = true;
                    } else {
                        dfs(j, i);
                        if (found) return;
                        if (has_forbidden[j]) {
                            has_forbidden[i] = true;
                            setmax(snd_forbidden, dp[j]);
                            if (max_forbidden < snd_forbidden) swap(max_forbidden, snd_forbidden);
                        } else {
                            sum_free += dp[j];
                        }
                    }
                }
                dp[i] = 1 + sum_free + max_forbidden;
                if (k <= dp[i] + snd_forbidden) found = true;
            };
            dfs(root, -1);
            if (found) break;
        }
        (not found ? high : low) = mid;
    }
    cout << low << endl;
    return 0;
}