牛客练习赛56 B 小琛和他的学校

题目链接：https://ac.nowcoder.com/acm/contest/3566/B

思路：一条路可把图分为左右两部分。

l_ci, l_peo, r_ci, r_peo, w 分别为左边城市数和人数，右边城市数和人数，该路的费用。

我们知道，左边的人要去右边的r_ci个城市，右边的人要去左边的l_ci个城市，

那么该路的费用就是 cost = 2*w*(l_ci*r_peo + r_ci*l_peo);（注意来回）。

我们可以用拓扑排序来得出一条边的左右人数和城市情况。

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <queue>
using namespace std;
typedef long long ll;

const int N = (int)2e5+100;
struct node{
    int to,nxt;
}e[N<<1];
struct Info{
    int ci,peo;
    Info(){ci = peo = 0;}
}info[N];
int A[N],W[N],du[N],head[N];//人数，路费,度数
bool vis[N];
int n,u,v,w,tot;ll sum_peo;

inline void add(int u,int v){
    e[tot].to = v; e[tot].nxt = head[u]; head[u] = tot++;
    e[tot].to = u; e[tot].nxt = head[v]; head[v] = tot++;
}

void top_sort(){
    //度数为1的压入
    queue<int > que;
    for(int i = 1; i <= n; ++i){
        if(du[i] == 1){
            vis[i] = 1;
            que.push(i);
        }
    }

    while(!que.empty()){
        int now = que.front(); que.pop();
        //加上自己城市的情况
        info[now].ci += 1;
        info[now].peo += A[now];
        for(int o = head[now]; ~o; o = e[o].nxt){
            int to = e[o].to;
            if(!vis[to]){
                //传递自己城市的情况
                info[to].ci += info[now].ci;
                info[to].peo += info[now].peo;
                if(--du[to] == 1){
                    vis[to] = 1;
                    que.push(to);
                }
            }
        }
    }
}

void show_info(){
    for(int i = 1;i <= n; ++i){
        printf("now = %d city = %d people = %d
",i,info[i].ci,info[i].peo);
    }
}

inline ll cost(Info& u,int i){
    Info v;
    v.ci = n-u.ci;
    v.peo = sum_peo - u.peo;
  //  printf("u.c = %d v.c = %d  u.p = %d v.p = %d W = %d
",u.ci,v.ci,u.peo,v.peo,W[i]);
    return (ll)2*W[i]*((ll)u.ci*v.peo+(ll)v.ci*u.peo);
}

int main(){

    scanf("%d",&n);
    for(int i = 0; i <= n; ++i) head[i] = -1; tot = 0;
    for(int i = 1; i <= n; ++i){
        scanf("%d",A+i);
        sum_peo += A[i];
    }
    for(int i = 1; i < n; ++i){
        scanf("%d%d%d",&u,&v,W+i);
        add(u,v);
        ++du[u]; ++du[v];//无向边
    }
    top_sort();
    //show_info();
    for(int i = 0; i < 2*(n-1); i+=2){
        //我们总是选一条边城市数少的一边去得出另一边的人数和城市情况
        int x = info[e[i].to].ci < info[e[i^1].to].ci ? e[i].to : e[i^1].to;
     //   printf("city = %d
",x);
        printf("%lld
",cost(info[x],1+(i>>1)));
    }

    return 0;
}