2021-03-13

I Full Depth Morning Show(树形dp)

思路：
1：dfs1自底向上用下图①②的公式求出每个点的stu和stdu，最终求出stup[1]和stdu[1]
2：dfs2由上向下用公式2,从1开始更新sdu[i]和stdu[i]，递归的过程中用sdu[i]和stdu[i]算出ans[i]

#include <iostream>
#include <cstring>

using namespace std;

const int N = 100010, M = N *2;

typedef long long LL;

LL h[N], e[M], ne[M], w[M], ans[N], t[N], idx, n;
LL sum[N], siz[N], sdu[N], stdu[N];
LL SUM, num;

void add(int a, int b, int c) {
    e[idx] = b, w[idx] = c, ne[idx] = h[a], h[a] = idx ++;
}

void dfs1(int u, int father) {
    siz[u] = 1, sum[u] = t[u];
    for (int i = h[u]; i != -1; i = ne[i]) {
        int j = e[i];
        if (j == father)    continue;
        dfs1(j, u);
        siz[u] += siz[j], sum[u] += sum[j];
        sdu[u] += siz[j] * w[i] + sdu[j];
        stdu[u] += sum[j] * w[i] + stdu[j];

    }

}

void dfs2(int u, int father) {
    for (int i = h[u]; i != -1; i = ne[i]) {
        int j = e[i];
        if (j == father)    continue;
        sdu[j] = sdu[u] + (n - 2 * siz[j]) * w[i];
        stdu[j] = stdu[u] + (SUM - 2 * sum[j]) * w[i];
        ans[j] = sdu[j] * t[j] + stdu[j];
        dfs2(j, u);
    }
}
int main() {
    memset(h, -1, sizeof h);
    cin >> n;
    for (int i = 1; i <= n; i ++)
        cin >> t[i], SUM += t[i];

    for (int i = 1; i < n; i ++) {
        int a, b, c;
        cin >> a >> b >> c;
        add(a, b, c), add(b, a, c);
    }

    dfs1(1, -1);
    ans[1] = t[1] * sdu[1] + stdu[1];
    dfs2(1, -1);
    for (int i = 1; i <= n; i ++)   cout << ans[i] << endl;


    return 0;

}