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

题目链接: 阔力梯的树

大致题意

如题所述.

想额外提一点的是, 题目中说这个节点的子树, 是需要计算当前节点的.

解题思路

dsu on tree

读完题, 我们分析一下这个题的特点, 需要统计每个节点所有子树的情况. 那不刚好符合dsu on tree能处理的问题吗? 考虑到信息维护, 我们直接把子树的信息维护到一个set中即可.

很快, 打完了就吃了T. 因为我每次求当前子树的答案时, 我都从头到尾跑了一遍set, 思考后发现, 还不如暴力方法来的快.

考虑如何优化 我们回想一下之前的dsu on tree题目, 我们需要在遍历当前节点所有轻儿子的时候, 就求出贡献, 这样才能保证O(nlogn)的复杂度.

在遍历轻儿子的函数中(下面代码为fact函数), 我们的操作就是, 把x节点加入set中. 我们发现, 其实插入x后贡献的变化情况我们可以O(1)求出.

​ ①插入x后, x位于set最开头/结尾. 此时相当于在原有贡献的基础上, 直接加上x产生的新贡献即可.
​ ②插入x后, x位于两个已有编号中间, 此时我们需要减去原本两个编号产生的贡献, 同时在加上x和两者产生的新贡献. (若其前一个数为prev, 后一个为next, 则原本的贡献是 prev和next 产生的. 插入后变为 prev和x, x和next共同产生.)

AC代码

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 1; i <= (n); ++i)
using namespace std;
typedef long long ll;
const int N = 1E5 + 10;
vector<int> edge[N];
int son[N], sz[N];

ll res[N], sum = 0;
void dfs1(int x) {
	sz[x] = 1;
	for (auto& to : edge[x]) {
		dfs1(to);
		sz[x] += sz[to];
		if (sz[to] > sz[son[x]]) son[x] = to;
	}
}

set<int> st;
void fact(int x, int pson) {
	st.insert(x);
	auto index = st.lower_bound(x);
	if (index == st.begin()) { // next --> x next
		if (next(index) != st.end()) {
			ll temp = abs(*index - *next(index));
			sum += temp * temp;
		}
	}
	else if (next(index) == st.end()) { // prev --> prev x
		if (index != st.begin()) {
			ll temp = abs(*index - *prev(index));
			sum += temp * temp;
		}
	}
	else { // prev next --> prev x next
		ll temp = abs(*next(index) - *prev(index));
		sum -= temp * temp;

		temp = abs(*next(index) - *index);
		sum += temp * temp;

		temp = abs(*prev(index) - *index);
		sum += temp * temp;
	}

	for (auto& to : edge[x]) {
		if (to == pson) continue;
		fact(to, pson);
	}
}
void dfs2(int x, int tp) {
	for (auto& to : edge[x]) {
		if (to == son[x]) continue;
		dfs2(to, 0);
	}

	if (son[x]) dfs2(son[x], 1);

	fact(x, son[x]);
	res[x] = sum;

	if (!tp) st.clear(), sum = 0;
}

int main()
{
	int n; cin >> n;
	rep(i, n - 1) {
		int x; scanf("%d", &x);
		edge[x].push_back(i + 1);
	}

	dfs1(1);
	dfs2(1, 1);

	rep(i, n) printf("%lld\n", res[i]);
	return 0;
}

END