CSP202212-4---＞聚集方差

65分
#include
#include
#include
#include
using namespace std;
typedef long long LL;
const LL N=3e5+10;
vector<LL>son[N];
vector<LL>sets[N];
LL a[N];
LL n;
vector<LL>Find(LL x)
{
    vector<LL>ans;
    ans.push_back(x);
    if(son[x].size()==0)
    return ans;
    vector<LL>&t=son[x];
    vector<vector<LL>>s;
    for(LL i=0;i<t.size();i++)
    {
        s.push_back(Find(t[i]));
    }
    for(LL i=0;i<s.size();i++)
    {
        vector<LL>&k=s[i];
        for(LL j=0;j<k.size();j++)
        ans.push_back(k[j]);
    }
    sets[x]=ans;
    return ans;
}
//输入x,找到以它为根的树的全部编号
LL Count(vector<LL>v)
{
    LL ans=0;
    vector<LL>s;
    for(LL i=0;i<v.size();i++)
    {
        s.push_back(a[v[i]]);
    }
    if(s.size()==1)
    return 0;
    sort(s.begin(),s.end());
    LL size=s.size();
    for(LL i=0;i<size;i++)
    {
       LL tmp=0x3f3f3f3f3f3f3f3f;
       if(i+1>=0&&i+1<size)
       tmp=min((s[i+1]-s[i])*(s[i+1]-s[i]),tmp);
       if(i-1>=0&&i-1<size)
       tmp=min((s[i-1]-s[i])*(s[i-1]-s[i]),tmp);
       ans+=tmp;
       //cout<<"----"<
    }
    return ans;
}
//给定一个集合,求解集合中元素的G值
int main()
{
    scanf("%lld",&n);
    for(LL i=2;i<=n;i++)
    {
        LL tmp;
        scanf("%lld",&tmp);
        son[tmp].push_back(i);
    }
    //build tree
    for(LL i=1;i<=n;i++)
    scanf("%lld",&a[i]);
    Find(1);
    for(LL i=1;i<=n;i++)
    {
        cout<<Count(sets[i])<<endl;
    }
    //cout<
}

满分代码：dsu on tree+multiset
#include
#include
#include
#include
#include
using namespace std;
typedef long long LL;
const LL N=3e5+10,M=2*N;
LL h[N],e[M],ne[M],idx;
LL n;
LL a[N];//记录每个点的值是多少
void add(LL a,LL b)
{
    e[idx]=b,ne[idx]=h[a],h[a]=idx++;
}
LL sz[N],son[N];//计算每个点的儿子
multiset<LL>store;
LL sum;
LL ans[N];
//记录每个点的答案是多少
void dfs_son(LL u,LL fa)
{
    sz[u]=1;
    LL big_son=0;
    LL maxv=0;
    for(LL i=h[u];~i;i=ne[i])
    {
        LL t=e[i];
        if(t==fa)
        continue;
        //不往上搜索
        dfs_son(t,u);
        sz[u]+=sz[t];
        if(sz[t]>maxv)
        {
            maxv=sz[t];
            big_son=t;
        }
    }
    son[u]=big_son;
}
LL count(LL x1,LL x2,LL x3)
{
    if(abs(x2)==0x3f3f3f3f)
    return 0;
    //null
    if(abs(x1)==0x3f3f3f3f&&abs(x3)==0x3f3f3f3f)
    {
        return 0;
    }
    else if(abs(x1)==0x3f3f3f3f&&abs(x3)!=0x3f3f3f3f)
    {
        return (x3-x2)*(x3-x2);
    }
    else if(abs(x3)==0x3f3f3f3f&&abs(x1)!=0x3f3f3f3f)
    {
        return (x1-x2)*(x1-x2);
    }
    else if(abs(x1)!=0x3f3f3f3f&&abs(x3)!=0x3f3f3f3f)
    {
        return min((x3-x2)*(x3-x2),(x1-x2)*(x1-x2));
    }
}
void addOne(LL val)
{
    auto t=store.lower_bound(val);
    --t,--t;
    //找到前驱
    LL be[6],af[6];
    af[1]=*t,t++;
    af[2]=*t,t++;
    af[3]=val;
    af[4]=*t,t++;
    af[5]=*t,t++;
    be[1]=af[1];
    be[2]=af[2];
    be[3]=af[4];
    be[4]=af[5];
    sum-=count(be[1],be[2],be[3]);
    sum-=count(be[2],be[3],be[4]);
    sum+=count(af[1],af[2],af[3]);
    sum+=count(af[2],af[3],af[4]);
    sum+=count(af[3],af[4],af[5]);
    store.insert(val);
}
void addAll(LL u,LL fa)
{
    addOne(a[u]);
    for(LL i=h[u];~i;i=ne[i])
    {
        LL t=e[i];
        if(t==fa)
        continue;
        addAll(t,u);
    }
}
void init()
{
    store.insert(-0x3f3f3f3f);
    store.insert(-0x3f3f3f3f);
    store.insert(0x3f3f3f3f);
    store.insert(0x3f3f3f3f);
}
void clear()
{
    store.erase(store.begin(),store.end());
    init();
    sum=0;
}
void func(LL u,LL fa,LL type)
{
    for(LL i=h[u];~i;i=ne[i])
    {
        LL t=e[i];
        if(t==fa||t==son[u])
        continue;
        //不往上搜索,并且最后搜索重儿子
        func(t,u,0);
    }
    if(son[u])
    func(son[u],u,1);
    //如果不是叶结点,就一定有重儿子,搜索之,并保存结果
    addOne(a[u]);
    for(LL i=h[u];~i;i=ne[i])
    {
        LL t=e[i];
        if(t==fa||t==son[u])
        continue;
        addAll(t,u);
    }
    //往set中加入全部的轻儿子的值
    ans[u]=sum;
    if(!type)
    {
        clear();
    }
    //需要自我清理的情况
}
//该函数可以求解出u的答案
//type==1:不用自我清理
//type==0:需要自我清理
int main()
{
    init();
    memset(h,-1,sizeof h);
    //清空邻接表的表头
    scanf("%lld",&n);
    for(LL i=2;i<=n;i++)
    {
        LL a=i,b;
        scanf("%lld",&b);
        add(a,b),add(b,a);
    }
    //建树
    for(LL i=1;i<=n;i++)
    scanf("%lld",&a[i]);
    //输入每个点的权值
    dfs_son(1,-1);
    //找到每一个点的重儿子
    func(1,-1,1);
    for(LL i=1;i<=n;i++)
    cout<<ans[i]<<endl;
}