2020 CCPC Wannafly Winter Camp 阔力梯的树 （树上启发式合并+set）

原题链接

分析：题中维护的是子树信息，自然会想到树启。

考虑到最朴素的做法，把每颗子树内的节点都取出来，然后sort一下，求个答案，时间复杂度在n * log(n) * n *log(n),因此肯定会超时，所以思考优化的方案。

每次加入一个节点，对答案的影响其实是可以O（1）处理出来的，然后加上找到这个点所花的log(n)的复杂度，其实最终的复杂度也只有n * log(n) * log(n),是符合要求的。

这里我用set来维护子树有序，然后用自带的二分去找位置， 注意一下边界的细节。

#include <iostream>
#include <algorithm>
#include <cstdio>
#include <cstring>
#include <queue>
#include <stack>
#include <cmath>
#include <bitset>
#include <map>
#include <set>
#include <time.h>

//#define ACM_LOCAL
using namespace std;
typedef long long ll;
typedef pair<int, int> PII;
const int N = 1e5 + 5, M = 5e5 + 5, INF = 0x3f3f3f3f;
const int MOD = 1e9 + 7;

int sz[N], son[N], h[N], tot, col[N], cnt[N], vis[N], mx;
ll sum, ans[N];
set<int> st;
struct Edge {
    int to, next;
}e[M];

void add(int u, int v) {
    e[tot].to = v;
    e[tot].next = h[u];
    h[u] = tot++;
}

void dfs(int u, int fa) {
    sz[u] = 1;
    for (int i = h[u]; ~i; i = e[i].next) {
        int v = e[i].to;
        if (v == fa) continue;
        dfs(v, u);
        sz[u] += sz[v];
        if (!son[u] || sz[v] > sz[son[u]])
            son[u] = v;
    }
}

void add(int u) {
    auto it = st.lower_bound(u), be = st.begin(), ed = st.end();
    if (be != ed) {
        auto edd = ed;
        --edd;
        if (it == be) sum += 1ll*(*be - u) * (*be - u);
        else if (it == ed) sum += 1ll*(u - *edd) * (u - *edd);
        else {
            auto t = it;
            --t;
            sum -= 1ll*(*it - *t) * (*it - *t);
            sum += 1ll*(u - *t) * (u - *t);
            sum += 1ll*(*it - u) * (*it - u);
        }
    } else {
        sum = 0;
    }
}

void del(int u) {
    auto it = st.lower_bound(u), be = st.begin(), ed = st.end();
    if (be != ed) {
        auto edd = ed;
        --edd;
        if (it == be) sum -= 1ll*(*be - u) * (*be - u);
        else if (it == ed) sum -= 1ll*(u-*edd) * (u - *edd);
        else {
            auto t = it;
            --t;
            sum += 1ll*(*it - *t) * (*it - *t);
            sum -= 1ll*(u - *t) * (u - *t);
            sum -= 1ll*(*it - u) * (*it - u);
        }
    } else {
        sum = 0;
    }
}

void count(int u, int fa, int k) {
    //统计子树信息
    if (k == 1) {
        add(u);
        st.insert(u);
    } else {
        st.erase(u);
        del(u);
    }
    //
    for (int i = h[u]; ~i; i = e[i].next) {
        int v = e[i].to;
        if (v == fa || vis[v]) continue;
        count(v, u, k);
    }
}

void dsu(int u, int fa, int keep) {
    for (int i = h[u]; ~i; i = e[i].next) {
        int v = e[i].to;
        if (v == fa || son[u] == v) continue;
        dsu(v, u, 0);//查询轻儿子
    }
    if (son[u]) dsu(son[u], u, 1), vis[son[u]] = 1;//查询重儿子
    count(u, fa, 1);//统计子树信息
    ans[u] = sum;
    //统计答案
    if (son[u]) vis[son[u]] = 0;
    if (!keep) count(u, fa, -1), sum = 0;//撤回信息
}
int n;
void solve() {
    cin >> n;
    memset(h, -1, sizeof h);
    for (int i = 2; i <= n; i++) {
        int x; cin >> x;
        add(i, x);
        add(x, i);
    }
    dfs(1, 0);
    dsu(1, 0, 0);
    for (int i = 1; i <= n; i++)
        printf("%lld\n", ans[i]);
}

signed main() {
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
#ifdef ACM_LOCAL
    freopen("input", "r", stdin);
    freopen("output", "w", stdout);
#endif
    solve();
    return 0;
}