【题解】P5533 [CCO 2019] Winter Driving

P5533 [CCO 2019] Winter Driving

题意

给定一棵 \(n\) 个节点的树，每个节点有权值 \(a_i\)，现在请你给每一条边定向。

在定向之后，如果 \(u\) 能到达与自己不同的节点 \(v\)，则会产生 \(a_u\times a_v\) 的贡献。

特殊地，每个节点 \(u\) 都能到达自己，会产生 \(a_u\times (a_u-1)\) 的贡献。

所有满足上述贡献机制的点对产生的贡献之和为这棵树的总贡献。

请你找出使得总贡献最大的定向方案，并输出此时总贡献的最大值。

特殊地，每个点的度数不超过 \(36\)。

题解

知识点：折半搜索，树的重心。

挺好的题目，正好复习了一些树的性质。

先发现钦定一个点为根后，在最优方案下，根的每个子树内部每一条边都是同向的，也就是说子树要么是内向的要么是外向的，用调整法可以证明。

贡献分为三部分，内/外向子树对根的贡献、内向子树外向子树之间的贡献、以及子树内部贡献。

子树内部贡献与结构是内向树还是外向树无关，可以提前算出。

边方向向根的子树对根的贡献也很好算。

内/外向子树之间的贡献相当于他们两个的各自的权值和（即 \(a_i\) 和）相乘。

内/外向子树权值和的加和是定值，贡献是其乘积，显然内向个外向子树权值和一定是各接近总权值和一半更优。

而题目规定每个节点度数 \(d \le 36\)，几乎明示着用 meet in middle 来划分所选的根的儿子的方向是向根还是向叶。

考虑枚举根，如果当前根不是重心的话，肯定存在一个子树权值和大于总权值和的一半，直接选那个就行了。

所以只会在重心为根的时候做 meet in middle，而重心至多有 \(2\) 个。

时间复杂度是 \(O(n+2^{\frac{d}{2}})\)。

#include<bits/stdc++.h>
using namespace std;

#define rep(i,l,r) for(int i=(l);i<=(r);++i)
#define per(i,l,r) for(int i=(r);i>=(l);--i)
#define pr pair<int,int>
#define fi first
#define se second
#define pb push_back
#define all(x) (x).begin(),(x).end()
// #define sz(x) (x).size()
#define bg(x) (x).begin()
#define ed(x) (x).end()

#define N 202506
#define int long long

int n,a[N],fa[N],ans=0,siz[N];
vector<int>e[N];
unordered_map<int,int>mp[N];

inline void dfs(int k){
    siz[k]=a[k];

    for(int x:e[k]){
        if(x==fa[k]){
            continue;
        }
        dfs(x);
        siz[k]+=siz[x];
    }
}

inline int sz(int k,int p){
    if(fa[k]==p){
        return siz[k];
    }

    return siz[1]-siz[p];
}

inline int calc(int k,int p){
    if(mp[k].count(p)){
        return mp[k][p];
    }

    int ans=0;

    for(int x:e[k]){
        if(x==p){
            continue;
        }

        ans+=calc(x,k)+sz(x,k)*a[k];
    }

    return mp[k][p]=ans;
}

signed main(){
    // freopen(".in","r",stdin);
    // freopen(".out","w",stdout);
    ios::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);

    cin>>n;

    rep(i,1,n){
        cin>>a[i];
    }

    rep(i,2,n){
        cin>>fa[i];
        e[fa[i]].pb(i);
        e[i].pb(fa[i]);
    }

    dfs(1);

    rep(u,1,n){
        int sum=siz[1]-a[u],res=0,mx=0,add=0;

        vector<int>v;

        for(int x:e[u]){
            res+=calc(x,u)+sz(x,u)*a[u];

            if(sz(x,u)>mx){
                mx=sz(x,u);
            }

            v.pb(sz(x,u));
        }

        if(mx>sum/2){
            add=mx*(sum-mx);
        }
        else{
            int len=v.size(),mid=len/2;
            int tot1=(1<<mid)-1,tot2=(1<<(len-mid))-1;
            vector<int>w;

            rep(s,0,tot1){
                int tmp=0;

                rep(i,0,mid-1){
                    if(!(s>>i&1)){
                        continue;
                    }
                    tmp+=v[i];
                }

                w.pb(tmp);
            }

            sort(all(w));

            rep(s,0,tot2){
                int tmp=0;

                rep(i,mid,len-1){
                    int b=i-mid;
                    if(!(s>>b&1)){
                        continue;
                    }
                    tmp+=v[i];
                }

                auto it=lower_bound(all(w),sum/2-tmp);

                if(it==ed(w)){
                    continue;
                }

                // res=max(res,(tmp+*it)*(sum-(tmp+*it)));
                add=max(add,(tmp+*it)*(sum-(tmp+*it)));
            }
        }

        ans=max(ans,res+add);
    }

    rep(i,1,n){
        ans+=a[i]*(a[i]-1);
    }

    cout<<ans;

    return 0;
}