Official

G - Tree Inversion Editorial by en_translator


First, we consider which \(u\) a pair \((v,w)\ (v\lt w)\) contributes to.

As \(v\lt w\), it should hold that \(w\) is on the \(u-v\) path, which holds if and only if, in the forest obtained by removing \(w\), \(v\) and \(u\) belongs to different connected components.

Taking an arbitrary root to regard it as a rooted tree, the set of such vertices \(u\) coincides with one of:

  • a subtree, or
  • all vertices but ones in a subtree.

For a fixed subtree, we count the pairs \((v,w)\) that contributes in each way.

1. A subtree

Fix a subtree consisting of a vertex \(p\) and its descendants.

When a pair \((v,w)\ (v\lt w)\) satisfies the following conditions, the set of vertices \(u\) that \((v,w)\) contributes to coincides with the subtree:

  • \(w=p\);
  • \(v\) is not in this subtree.

The number of such pairs \((v,w)\) can be found as \(w-1-(\) the number of the subtree’s vertices \(x\) such that \(x\lt w)\).

2. All but those in a subtree.

Fix a subtree consisting of a vertex \(p\) and its descendants.

When a pair \((v,w)\ (v\lt w)\) satisfies the following conditions, the set of vertices \(u\) that \((v,w)\) contributes to coincides with the entire vertices except for those in the subtree:

  • \(w\) is the parent of \(p\);
  • \(v\) is included in this subtree.

The number of such pairs \((v,w)\) can be found as the number of the subtree’s vertices \(x\) such that \(x\lt w\).

Therefore, the contributions can be computed fast if one can find the number of the subtree’s vertices such that \(x\lt q\) for any subtree and an integer \(q\).
By plotting the preorder of the DFS (Depth-First Search) and the values, it turns out to be a counting query on a rectangular region. By performing 2D scanning in the preorder of the DFS, it is boiled down to the prefix sum of a sequence.
With Fenwick Tree or a segment tree, the queries can be processed fast.

In order to find the answer from the contributions obtained, one has to add a constant to all vertices in a subtree, or to all vertices but in a subtree. This can be achieved with a cumulative sum trick on a tree, or with Euler tour.

The following is sample code. This implementation adopts Fenwick tree for the prefix sum part, and a cumulative sum trick on a tree for batch-adding to a subtree.

#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] := values added to all vertices in a subtree rooted at i
    vector<unsigned long> imos(N);

    // ∑_{i<x} prefix_sum[i] := The number of vertices visited so far, whose numbers is no greater than x
    atcoder::fenwick_tree<unsigned long> prefix_sum(N);

    // Finds the contribution of each subtree
    // O(N log N) time
    [dfs{[&edges, &imos, &prefix_sum](const auto &f, unsigned now, unsigned prev) -> void {
        // Record the prefix sum when entering this vertex
        auto child_in{prefix_sum.sum(0, now)}, not_child_in{prefix_sum.sum(0, prev)};
        prefix_sum.add(now, 1UL); // We have visited this vertex, so prefix_sum[now] += 1

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

        // Find the prefix sum when exiting
        auto child_out{prefix_sum.sum(0, now)}, not_child_out{prefix_sum.sum(0, prev)};

        // Add `now` - (#vertices less than `now` in the subtree) to the subtree
        // Add (#vertices less than `prev` in the subtree) to the complement of the subtree
        auto child_add{now + child_in - child_out}, not_child_add{not_child_out - not_child_in};

        // Cumulative sum trick (add to the complement = add globally, then subtract from the subtree)
        imos[now] += child_add - not_child_add;
        imos[0] += not_child_add;
    }}] {
        dfs(dfs, 0, 0);
    }();

    // Complete the cumulative sum trick
    [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);
    }();

    // Print the answer
    for (const auto ans : imos)
        cout << ans << " ";
    cout << endl;

    return 0;
}

posted:
last update: