清北学堂模拟赛d4t5 b

分析：一眼树形dp题，就是不会写QAQ.树形dp嘛，定义状态肯定有一维是以i为根的子树,其实这道题只需要这一维就可以了.设f[i]为以i为根的子树中的权值和.先处理子树内部的情况，用一个数组son[i]表示以i为根的子树中,i能走到的节点个数，可以利用son数组和当前点的权值来更新f数组.

处理了每个子树内部的情况，接下来就要合并它们，将每一个根节点作为中间点，算一下中间点权值的贡献，利用乘法原理算出有多少对点对经过中间点，乘一下就ok了.

树形dp的基本状态定义要熟记，有些题目子树内部是互相独立的，可以在子树里面单独计算，最后再合并一下.

#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>

using namespace std;

const int maxn = ;

int n, a[maxn], head[maxn], to[maxn * ], nextt[maxn * ], tot = , w[maxn * ];
long long ans, f[maxn], son[maxn];

void add(int x, int y, int z)
{
    w[tot] = z;
    to[tot] = y;
    nextt[tot] = head[x];
    head[x] = tot++;
}

void dfs(int u, int fa, int col)
{
    long long res = ;
    f[u] = a[u];
    son[u] = ;
    bool flag = ;
    for (int i = head[u]; i; i = nextt[i])
    {
        int v = to[i];
        if (v == fa)
            continue;
            dfs(v, u, w[i]);
            if (col != w[i])
            {
                flag = ;
                son[u] += son[v];
                f[u] += son[v] * a[u] + f[v];
            }
            res += son[v] * a[u] + f[v];
    }
    ans += res;
    if (flag)
        return;
    for (int i = head[u]; i; i = nextt[i])
    {
        int v1 = to[i];
        if (v1 != fa)
            for (int j = i; j; j = nextt[j]) //防止重复统计，所以j=i而不是j=head[u]
            {
                int v2 = to[j];
                if (v2 != fa && w[i] != w[j])
                    ans += son[v1] * f[v2] + son[v2] * f[v1] + a[u] * son[v1] * son[v2];
            }
    }
}

int main()
{
    scanf("%d", &n);
    for (int i = ; i <= n; i++)
        scanf("%d", &a[i]);
    for (int i = ; i < n; i++)
    {
        int x, y, z;
        scanf("%d%d%d", &x, &y, &z);
        add(x, y, z);
        add(y, x, z);
    }
    dfs(, , );
    printf("%lld\n", ans);

    return ;
}