2020 ccpcwannafly day2 E题

阔力梯的树

题目的主要思想是dsu on tree（树上启发式合并），可以解决某些树上离线问题。
大致过程有三步：
1.dfs预处理出每个每个节点子树的大小，dfs序（有时候还有对应dfs序的节点编号），重儿子。
2.第二遍dfs，先遍历轻儿子，并且不保留答案，然后遍历重儿子，保留，然后把所有的轻儿子的信息加入
3.如果不保留的话就整体清空

在这个题目中，我们需要把树上每个节点的子树编号拿出来，排序后两两求差后的平方和。
这里，我们可以用set维护子树的序列，逐个加入的时候把它和相邻的元素拿出来对答案加上贡献。

#include <iostream>
#include <string>
#include <cstring>
#include <algorithm>
#include <set>
using namespace std;
const int N = 100010;
typedef long long ll;
int son[N], sz[N];//son里面存的是重儿子
ll res[N];
set<int>s;
int h[N], e[N], ne[N], idx, n, cnt, id[N], l[N];

void add(int a, int b) {
	e[idx] = b, ne[idx] = h[a], h[a] = idx++;
}

void dfs1(int x) {
	sz[x] = 1;
	l[x] = ++cnt;//遍历得到dfs序
	id[cnt] = x;//对应dfs序的点的编号
	for (int i = h[x]; i != -1; i = ne[i]) {
		int y = e[i];
		dfs1(y);
		sz[x] += sz[y];//求出每个点的size
		if (sz[y] > sz[son[x]]) {
			son[x] = y;//启发式合并
		}
	}
}

void modify(int x, int val) {
	set<int>::iterator it = s.lower_bound(val);
	if (s.size() == 0) {
		s.insert(val);
		return;
	}
	if (it == s.begin()) {//最小的比它大
		ll tmp = (*it) - val;
		res[x] += tmp * tmp;//贡献
	}
	else if (it == s.end()) {
		it--;
		ll tmp = val - (*it);
		res[x] += tmp * tmp;//贡献
	}
	else {
		ll tmp = (*it) - (*(--it));
		res[x] -= tmp * tmp;
		tmp = val - (*it);
		res[x] += tmp * tmp;
		tmp = (*(++it)) - val;
		res[x] += tmp * tmp;
	}
	s.insert(val);
}

void dfs2(int x,bool flag) {
	for (int i = h[x]; i != -1; i = ne[i]) {//先遍历一遍轻儿子，并且不保留
		int y = e[i];
		if (y == son[x])continue;
		dfs2(y, false);
	}
	
	if (son[x])dfs2(son[x], true);
	res[x] = res[son[x]];
	for (int i = h[x]; i != -1; i = ne[i]) {//插入轻儿子，保留
		int y = e[i];
		if (y == son[x])continue;
		for (int j = l[y]; j <= l[y] + sz[y] - 1; j++) {//把加入控制在子树内
			//cout << x << ' ' << id[j] << ' ' << endl;
			//添加，计算答案
			modify(x, id[j]);
		}

	}
	
	modify(x, x);
	if (!flag) {
		s.clear();
	}
}

int main() {
	cin >> n;
	memset(h, -1, sizeof h);
	for (int i = 2; i <= n; i++) {
		int x;
		cin >> x;
		add(x, i);
	}

	dfs1(1);
	dfs2(1, false);
	for (int i = 1; i <= n; i++) {
		cout << res[i] << endl;
	}
	return 0;
}

dfs1是预处理，dfs2是第二部分，modify是计算在加入轻儿子时对答案产生的贡献。贡献有三种情况，跟