Cut and Add

Cut and Add

差分 dfn序

#include <bits/stdc++.h>
using namespace std;
using ll = long long;

int n;
vector<vector<int>> g;
vector<int> pa, sz;
vector<int> tim;

void predfs(int u, int p) {
    tim.push_back(u);
    for(int v : g[u]) {
        if(v == p) continue;
        pa[v] = u;
        predfs(v, u);
        sz[u] += sz[v];
    }
}

void solve() {
    cin >> n;
    g.assign(n + 1, vector<int>());
    pa.assign(n + 1, 0); pa[1] = -1;
    sz.assign(n + 1, 1);

    vector<pair<int, int>> edge(n + 1);
    for(int i = 1; i < n; i ++) {
        int a, b; cin >> a >> b;
        edge[i] = {a, b};
        g[a].push_back(b);
        g[b].push_back(a);
    }

    predfs(1, -1);
    vector<int> id(n + 1);
    for(int i = 1; i <= n; i ++) {
        id[tim[i - 1]] = i;
    }
    vector<ll> ans(n + 1, 0), dif(n + 2, 0);
    for(int i = 1; i <= n; i ++) {
        auto [x, y] = edge[i];
        if(pa[x] != y) swap(x, y);
        if(sz[id[x]] <= n / 2) {
            int r = id[x];
            int l = r - sz[x] + 1;
            dif[1] += i, dif[n + 1] -= i;
            dif[r + 1] += i; dif[l] -= i;
        }
        else {
            int r = id[x];
            int l = r - sz[x] + 1;
            dif[l] += i, dif[r] -= i;
        }
    }
    for(int i = 1; i <= n; i ++) ans[i] = ans[i - 1] + dif[i];
    for(int i = 1; i <= n; i ++) cout << ans[i] << " \n"[i == n];
}

int main() {
    ios::sync_with_stdio(false); cin.tie(0); cout.tie(0);
    int _ = 1; //cin >> _;
    while(_--) solve();
    return 0;
}