gym 100496 House of Representatives(树形dp)

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;

typedef long long ll;
#define N 1111111

ll dp
;
int a
;
ll sum
,ans
;

struct EDGE {
    int v, next;
    ll w;
} edge
;

int head
, e;
int n, k, vis
;

void init() {
    memset(vis, 0, sizeof(vis));
    memset(head, -1, sizeof(head));
    e = 0;
}

void add(int u, int v, int w) {
    edge[e].v = v;
    edge[e].w = w;
    edge[e].next = head[u];
    head[u] = e++;
}

void dfs_dis(int u,int fa) {
    //printf("%d %d\n",u,fa);
    sum[u] = a[u];
    for(int i = head[u]; i != -1; i = edge[i].next) {
        int v = edge[i].v;
        if(v != fa &&!vis[v]) {
            vis[v]=1;
            dfs_dis(v, u);
            sum[u] += sum[v];
            dp[u] += dp[v]+sum[v]*edge[i].w;
        }
    }
}

void dfs(int u, int fa,ll val) {
    ans[u] =  dp[u] + val;
    for(int i = head[u]; i != -1; i = edge[i].next) {
        int v = edge[i].v;
        if(v != fa ) {
            ll tmp =  val;
            tmp += (dp[u]-dp[v]);
            tmp += ((sum[1]-sum[v])*edge[i].w-sum[v]*edge[i].w);
            dfs(v, u,tmp);
        }
    }
}

int main () {
    freopen("house.in","r",stdin);
    freopen("house.out","w",stdout);
    scanf("%d",&n);
    for(int i=1; i<=n; i++) {
        scanf("%d",&a[i]);
    }
    int c,b;
    ll w;
    init();
    for(int i=1; i<n; i++) {
        scanf("%d%d%I64d",&c,&b,&w);
        add(c,b,w);
        add(b,c,w);
    }

    dfs_dis(1,-1);

    dfs(1,-1,0);

    int k = 1;

    for(int i=1; i<=n; i++) {
        if(ans[k]>ans[i]) {
            k=i;
        }
    }
    printf("%d %I64d\n",k,ans[k]);
}