2020/11/06 模拟赛 T1

Description

有一$n$个节点的树，每个节点有一点权$a_i$，定义$dist(x,y)$为$x$到$y$的边数。选取一点$v$，使$\sum_{i=1}^n dist(i,v) \times a_i + dist(i,v)^2 \times b_i$最大

Solution

两次DFS，一次DFS求子树内，第二次考虑子树外

#include<iostream>
#include<cstdio>
#include<cmath>
using namespace std;
int n,tot,head[300005];
long long a[300005],b[300005],sa[300005],fa[300005],sb[300005],fb[300005],g[300005],ans;
struct Edge
{
    int to,nxt;
}edge[600005];
inline int read()
{
    int w=0,f=1;
    char ch=0;
    while(ch<'0'||ch>'9'){if(ch=='-') f=-1; ch=getchar();}
    while(ch>='0'&&ch<='9')w=(w<<1)+(w<<3)+ch-'0',ch=getchar();
    return w*f;
}
void dfs1(int k,int f)
{
    sa[k]=a[k],sb[k]=b[k];
    for(int i=head[k];i;i=edge[i].nxt)
    {
        int v=edge[i].to;
        if(v!=f)
        {
            dfs1(v,k);
            sa[k]+=sa[v],sb[k]+=sb[v];
            fa[k]+=fa[v]+sa[v],fb[k]+=fb[v]+sb[v];
            g[k]+=g[v]+2*fb[v]+sb[v];
        }
    }
}
void dfs2(int k,int f)
{
    if(k!=1)
    {
        g[k]+=(g[f]-g[k]-2*fb[k]-sb[k])+2*(fb[f]-fb[k]-sb[k])+sb[1]-sb[k];
        fa[k]+=fa[f]-fa[k]-sa[k]+sa[1]-sa[k],fb[k]+=fb[f]-fb[k]-sb[k]+sb[1]-sb[k];
    }
    for(int i=head[k];i;i=edge[i].nxt)
    {
        int v=edge[i].to;
        if(v!=f) dfs2(v,k);
    }
}
int main()
{
    n=read();
    for(int i=1;i<=n;i++) a[i]=read();
    for(int i=1;i<=n;i++) b[i]=read();
    for(int i=1;i<n;i++)
    {
        int u=read(),v=read();
        edge[++tot]=(Edge){v,head[u]},head[u]=tot;
        edge[++tot]=(Edge){u,head[v]},head[v]=tot;
    }
    dfs1(1,0);
    dfs2(1,0);
    for(int i=1;i<=n;i++) ans=max(ans,fa[i]+g[i]);
    printf("%lld\n",ans);
    return 0;
}