Birjik and Nicole's Tree Game | HourRank 20

復習用コメント付きコード

/*
思いつき方
頂点を黒く塗る
その黒い頂点を含む部分木は？
根まで行く距離が長過ぎる
HL分解する
*/

/*
HL分解
Heavy-Light Decomposition
木の構築 O(|V|+|E|)
あらかじめ与えられた木について、頂点u,v間のパスをセグメント木などで扱える半開区間の集合に変換する。区間の数はO(log(|V|))個。
ひとつの区間は連結した頂点たちにIDをふったもので、頂点の深さは全て異なる。その区間内では、(深さの値の差)=(IDの値の差) になっている．
1クエリ O(log(|V|))
スタックオーバーフロー回避版
*/
class HLDecomposition {
    int cur = 0;
    void bfs(const vector<vector<int> > &G) {
        queue<tuple<int, int, int> > que;
        vector<int> order;
        que.push(make_tuple(0, -1, 0));
        while (!que.empty()) {
            int u, pre, d;
            tie(u, pre, d) = que.front();
            que.pop();
            parent[u] = pre;
            depth[u] = d;
            order.push_back(u);
            for (int v : G[u])if (v != pre) {
                que.push(make_tuple(v, u, d + 1));
            }
        }
        reverse(order.begin(), order.end());
        for (int u : order) {
            sub[u] = 1;
            for (int v : G[u])if (v != parent[u])sub[u] += sub[v];
        }
    }

    void dfs_stk(const vector<vector<int> > & G) {
        stack<pair<int, int> > stk;
        stk.push(make_pair(0, 0));
        while (!stk.empty()) {
            int u, h, pre;
            tie(u, h) = stk.top();
            stk.pop();
            head[u] = h;
            id[u] = cur++;
            pre = parent[u];
            int heavy = -1, maxi = 0;
            for (int v : G[u]) {
                if (v != pre && maxi < sub[v]) {
                    maxi = sub[v];
                    heavy = v;
                }
            }
            for (int v : G[u]) {
                if (v != pre&&v != heavy) {
                    stk.push(make_pair(v, v));
                }
            }
            if (heavy != -1)stk.push(make_pair(heavy, h));
        }
    }
public:
    vector<int> parent, depth, sub, id, head;
    
    HLDecomposition(const vector<vector<int> > &G) {
        parent = depth = sub = id = head = vector<int>(G.size());
        bfs(G);
        dfs_stk(G);
    }

    /*
    パス(u, v)を半開区間の集合に変換する。
    O(log(|V|))
    */
    vector<pair<int, int> > goUpToLCA(int u, int v) {
        vector<pair<int, int> > res;
        while (true) {
            if (head[u] == head[v]) {
                if (depth[u] > depth[v])swap(u, v);
                res.emplace_back(id[u], id[v] + 1);
                break;
            }
            else {
                if (depth[head[u]] > depth[head[v]])swap(u, v);
                res.emplace_back(id[head[v]], id[v] + 1);
                v = parent[head[v]];
            }
        }
        return res;
    }
};

int N, Q, K;
vector<vi> G;
int vs[300001], ans[300001];

int main() {
    scanf("%d", &N);
    G.resize(N);
    rep(i, N - 1) {
        int u, v;
        scanf("%d%d", &u, &v);
        // 0-indexedとする。
        --u;
        --v;
        G[u].push_back(v);
        G[v].push_back(u);
    }

    HLDecomposition hld(G);
    // sgs[l]はlをheadとするセグメント[l, R)に対して
    // l<r<=R, r∈sgs[l]
    // を満たすような多重集合
    unordered_map<int, vi> sgs;

    scanf("%d", &Q);
    while (Q--) {
        sgs.clear();
        scanf("%d", &K);
        rep(i, K) {
            scanf("%d", &vs[i]);
            --vs[i];
            /*
            頂点vs[i]を含む部分木は
            vs[i], vs[i]の親、vs[i]の親の親、、、根
            つまりvs[i]の先祖すべてを根とする部分木すべて。
            これをすべて得るにはvs[i]から根までのパスを見るだけでいいが、
            O(|V|)かかる。
            かわりに、HL分解によってO(log|V|)個のパスを得る。
            */
            // log(|V|)個のパスに変換
            auto s = hld.goUpToLCA(0, vs[i]);
            each(p, s) {
                // ここが最も深い
                // O(log(N)ΣK)
                sgs[p.first].push_back(p.second);
            }
        }
        rep(i, K + 1)ans[i] = 0;

        each(p, sgs) {
            /*
            headのidをlとするようなパスについて考える。
            IDがlである頂点を根とする部分木は、黒い頂点をcn個含む。
            */
            int l = p.first, cn = sz(p.second);
            // rを昇順に見る。
            // より深い頂点へ向かっていく。
            sort(all(p.second));
            each(r, p.second) {
                /*
                r-1に対応する頂点が黒いので、
                rまで見たら黒色が1つ減る。
                [l, r)に対応するr-l個の頂点はcn個の黒い頂点を持つ。
                */
                ans[cn] += r - l;
                l = r;
                // 黒い頂点を1個減らす。(その黒い頂点よりも下へ行く)
                cn--;
            }
        }
        ans[0] = N;
        for (int i = 1; i <= K; ++i)ans[0] -= ans[i];
        rep(i, K + 1)printf("%d%c", ans[i], i != K ? ' ' : '\n');
    }
}