原题链接
分析:题中维护的是子树信息,自然会想到树启。
考虑到最朴素的做法,把每颗子树内的节点都取出来,然后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;
}