【面试向】hihoCoder 1994 树与落叶

题目链接

Implementation

    int n, q; scan(n,q);

    vi p(n + 1);
    vi nson(n + 1);

    up (i, 1, n) {
        scan(p[i]);
        nson[p[i]]++;
    }

    vi leaf;
    up (i, 1, n) {
        if (nson[i] == 0) leaf.pb(i);
    }

    vi cnt;
    cnt.pb(n);

    for (; !leaf.empty(); )  {
        cnt.pb(cnt.back() - SZ(leaf));
        vi tmp;
        FOR (x, leaf) {
            --nson[p[x]];
            if (nson[p[x]] == 0 && p[x] != 0) {
                tmp.pb(p[x]);
            }
        }
        leaf = std::move(tmp);
    }

    rep (q) {
        int x; scan(x);

        // 小于等于 x 的第一个数的下标
        auto i = lower_bound(all(cnt), x, greater<int>()) - cnt.begin();

        debug(i);
        if (i == SZ(cnt)) {
            println(SZ(cnt));
        }
        else if (i == 0) {
            println(i + 1);
        }
        else {
            int d1 = x - cnt[i];
            int d2 = cnt[i - 1] - x;
            println(d2 <= d1 ? i : i + 1);
        }
    }

Note

这道题的实现要点是如何维护每天的叶子集合。

key observation: 对于 $i \ge 3$，第 $i$ 天掉落的叶子一定是某些第 $i - 1$ 天掉落的叶子的父亲。

对应的代码是

        vi tmp;
        FOR (x, leaf) {
            --nson[p[x]];
            if (nson[p[x]] == 0 && p[x] != 0) {
                tmp.pb(p[x]);
            }
        }
        leaf = std::move(tmp);