公式

G - Tree Inversion 解説 by MMNMM


まず、組 \((v,w)\ (v\lt w)\) がどの \(u\) に寄与するかについて考えます。

\(v\lt w\) より、満たすべき条件は \(u,v\) パス上に \(w\) が存在することです。 これは「\(w\) を取り除いた森において \(v\) が属する連結成分に \(u\) が属していない」と言い換えられます。

適当な頂点を根とした根つき木を考えると、この条件を満たす \(u\) の集合は

  • ある部分木
  • ある部分木以外全部

のどちらかの形で表せます。

部分木を固定し、それぞれの形で寄与する組 \((v,w)\) について考えます。

1. 部分木

頂点 \(p\) およびその子孫からなる部分木について考えます。

\((v,w)\) が寄与する頂点 \(u\) がこの部分木と一致するような \((v,w)\ (v\lt w)\) は、次の条件を満たします。

  • \(w=p\)
  • \(v\) はこの部分木に含まれない

このような組 \((v,w)\) の個数は、\(w-1-(\)部分木の頂点のうち \(x\lt w\) を満たす \(x\) の個数\()\) として求められます。

2. 部分木以外全部

頂点 \(p\) およびその子孫からなる部分木について考えます。

\((v,w)\) が寄与する頂点 \(u\) がこの部分木以外全部と一致するような \((v,w)\ (v\lt w)\) は、次の条件を満たします。

  • \(w\) は \(p\) の親
  • \(v\) はこの部分木に含まれる

このような組 \((v,w)\) の個数は、\((\)部分木の頂点のうち \(x\lt w\) を満たす \(x\) の個数\()\) として求められます。

以上より、部分木と整数 \(q\) に対して \((\)部分木の頂点のうち \(x\lt q\) を満たすものの個数\()\) を高速に求めることができれば、それぞれの寄与を計算することができます。
これは、DFS の行きがけ順と値を平面にプロットすると矩形カウントクエリになります。 DFS の行きがけ順に平面走査を行うことで、列の prefix sum に帰着することができます。
Binary Indexed Tree やセグメント木などを用いることで、高速にクエリを処理することができます。

求めた寄与から答えを計算するには、部分木や部分木以外全部に含まれる頂点に対して定数を加えるような操作が必要です。 これは木上の imos 法や Euler tour などによって実現できます。

実装例は以下のようになります。 この実装では、prefix sum に Binary Indexed Tree を、部分木への加算に木上の imos 法を用いています。

#include <iostream>
#include <vector>
#include <atcoder/fenwicktree>

int main() {
    using namespace std;
    unsigned N;
    cin >> N;

    vector<vector<unsigned>> edges(N);
    for (unsigned i{1}, u, v; i < N; ++i) {
        cin >> u >> v;
        edges[--u].emplace_back(--v);
        edges[v].emplace_back(u);
    }

    // imos[i] := i に対応する部分木全体に足されている値
    vector<unsigned long> imos(N);

    // ∑_{i<x} prefix_sum[i] := これまでに訪れた頂点のうち番号が x 以下のものの個数
    atcoder::fenwick_tree<unsigned long> prefix_sum(N);

    // 各部分木への寄与を求める
    // O(N log N) 時間
    [dfs{[&edges, &imos, &prefix_sum](const auto &f, unsigned now, unsigned prev) -> void {
        // 入ってきたときの prefix sum の値を記録しておく
        auto child_in{prefix_sum.sum(0, now)}, not_child_in{prefix_sum.sum(0, prev)};
        prefix_sum.add(now, 1UL); // 訪れたので prefix_sum[now] += 1

        for (const auto next : edges[now])
            if (next != prev)
                f(f, next, now); // 再帰

        // 出るときの prefix sum の値を求める
        auto child_out{prefix_sum.sum(0, now)}, not_child_out{prefix_sum.sum(0, prev)};

        // 部分木に加えるのは now - (部分木の now 未満の頂点の個数)
        // 部分木以外全部に加えるのは (部分木の prev 未満の頂点の個数)
        auto child_add{now + child_in - child_out}, not_child_add{not_child_out - not_child_in};

        // imos 法をする (部分木以外全部に加える = 全体に加えて部分木から引く)
        imos[now] += child_add - not_child_add;
        imos[0] += not_child_add;
    }}] {
        dfs(dfs, 0, 0);
    }();

    // imos 法の後半
    [dfs{[&edges, &imos](const auto &f, unsigned now, unsigned prev, unsigned long value) -> void {
        imos[now] += value;
        for (const auto next : edges[now])
            if (next != prev)
                f(f, next, now, imos[now]);
    }}] {
        dfs(dfs, 0, 0, 0);
    }();

    // 答えを出力
    for (const auto ans : imos)
        cout << ans << " ";
    cout << endl;

    return 0;
}

投稿日時:
最終更新: