阔力梯的树 树上启发式合并

阔力梯的树

题意

一棵 $n(1\le n\le 10^5)$ 个结点的树，$1$ 是根节点。一个结点的结实程度是 ${\sum }_{i=1}^{k-1}(a_{i+1}-a_i{)}^2$ ，其中 $a_i$ 为子树的结点编号按照升序排序。求每个结点的结实程度。

解法

离线的子树问题，考虑树上启发式合并。

用一个 $set$ 来维护子树的所有结点，然后每次插入或者删除结点的时候，更新答案即可，注意分类讨论即可。

代码

#pragma region
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <vector>
using namespace std;
typedef long long ll;
#define tr t[root]
#define lson t[root << 1]
#define rson t[root << 1 | 1]
#define rep(i, a, n) for (int i = a; i <= n; ++i)
#define per(i, a, n) for (int i = n; i >= a; --i)
#pragma endregion
const int maxn = 1e5 + 5;
int n;
vector<int> g[maxn];
int sz[maxn], son[maxn];
int flag;
ll ans[maxn], sum;
set<ll> st;
void dfs1(int u) {
    sz[u] = 1;
    for (auto v : g[u]) {
        dfs1(v);
        sz[u] += sz[v];
        if (sz[v] > sz[son[u]]) son[u] = v;
    }
}
void count(int u, int f, int val) {
    if (val == 1) st.insert(u);
    auto it = st.find(u);
    auto en = st.end();
    --en;
    if (it == st.begin()) {
        ++it;
        if (it != st.end()) sum += val * (*it - u) * (*it - u);
    } else if (it == en) {
        --it;
        sum += val * (*en - *it) * (*en - *it);
    } else {
        auto p = --it;
        ++it;
        auto q = ++it;
        sum += val * ((*q - u) * (*q - u) + (u - *p) * (u - *p) - (*q - *p) * (*q - *p));
    }
    if (val == -1) st.erase(st.find(u));
    for (auto v : g[u]) {
        if (v == f || v == flag) continue;
        count(v, u, val);
    }
}
void dfs(int u, int f, bool keep) {
    for (auto v : g[u]) {
        if (v == f || v == son[u]) continue;
        dfs(v, u, 0);
    }
    if (son[u]) {
        dfs(son[u], u, 1);
        flag = son[u];
    }
    count(u, f, 1);
    ans[u] = sum;
    flag = 0;
    if (!keep) {
        count(u, f, -1);
    }
}
int main() {
    scanf("%d", &n);
    rep(i, 2, n) {
        int f;
        scanf("%d", &f);
        g[f].push_back(i);
    }
    dfs1(1);
    dfs(1, 0, 0);
    rep(i, 1, n) printf("%lld\n", ans[i]);
}