[洛谷-P1364] 医院设置

普通的floyd就不讲了，如果数据量到了1e5以上，这就是一道树的重心的变式，求带权的重心。或者说用树型dp或dfs来优化最小值的查找。最终时间复杂度 \(O(n)\) 。以下代码是第一篇题解的风格变化+注释。

#include <bits/stdc++.h>
#define int long long
using namespace std;
constexpr int maxn=1e4+10;// 题目数据小随便开
constexpr int maxm=2e4+10;
constexpr int INF=LONG_LONG_MAX>>1;

typedef struct edge
{
    int to,nex;
}edge;

edge edges[maxn];
int enums,heads[maxn];
int wi[maxn],siz[maxn],tal[maxn];
// siz[i]-以i为根的权值和
// tal[i]-以i为根的权值*路长的和
int ans=INF;

void add_edge(int x,int y)
{
    edges[++enums]={y,heads[x]};// 非常的压抑
    heads[x]=enums;
}

void dfs(int u,int fa,int dep)
{
    // 初始化siz和tal[1]
    siz[u]=wi[u];
    for(int i=heads[u];i;i=edges[i].nex)
    {
        if(edges[i].to!=fa)
        {
            dfs(edges[i].to,u,dep+1);//直到没有孩子
            siz[u]+=siz[edges[i].to];// 加上子树的权值
        }
    }
    tal[1]+=wi[u]*dep;
}

void dp(int u,int fa)
{
    // 从 u到 v，对于子树来说：距离-1，总贡献 -siz[v]。
    // 对其他点：距离+1，总贡献+siz[1]-siz[v]
    for(int i=heads[u];i;i=edges[i].nex)
    {
        if(edges[i].to!=fa)
        {
            tal[edges[i].to]=tal[u]-2*siz[edges[i].to]+siz[1];
            dp(edges[i].to,u);
        }
    }
    ans=min(ans,tal[u]);
}

signed main()
{
    int n;
    int a,b;
    scanf("%lld",&n);
    for(int i=1;i<=n;++i)
    {
        scanf("%lld%lld%lld",&wi[i],&a,&b);
        if(a)
        {
            add_edge(i,a);
            add_edge(a,i);
        }
        if(b)
        {
            add_edge(i,b);
            add_edge(b,i);
        }
    }
    dfs(1,0,0);
    dp(1,0);
    printf("%lld\n",ans);

    return 0;
}