比赛 2025.3.29 评测结果 AWWWWWWWWWAAAWAAAWWWWWWWW
题目名称 硝华流焰 最终得分 28
用户昵称 健康铀 运行时间 2.833 s
代码语言 C++ 内存使用 9.58 MiB
提交时间 2025-03-29 11:11:16
显示代码纯文本
#include <bits/stdc++.h>
using namespace std;
const int MOD = 1e9 + 7;
const int N = 1e5 + 5;

struct E {
    int v, w;
};
vector<E> g[N];
int c[N];
bool vs[N];
long long res = 0;
int sz(int u, int p) {
    int s = 1;
    for (E& e : g[u]) {
        if (e.v != p && !vs[e.v]) {
            s += sz(e.v, u);
        }
    }
    return s;
}
int find(int u, int p, int tot, int& ct) {
    int s = 1, mx = 0;
    for (E& e : g[u]) {
        if (e.v != p && !vs[e.v]) {
            int ch = find(e.v, u, tot, ct);
            s += ch;
            mx = max(mx, ch);
        }
    }
    mx = max(mx, tot - s);
    if (mx <= tot / 2) {
        ct = u;
    }
    return s;
}
void go(int u, int p, int m, long long sm, vector<long long>& cn, vector<long long>& sml) {
    cn[m]++;
    sml[m] = (sml[m] + sm) % MOD;
    for (E& e : g[u]) {
        if (e.v != p && !vs[e.v]) {
            int nm = m | (1 << c[e.v]);
            long long nsm = (sm + e.w) % MOD;
            go(e.v, u, nm, nsm, cn, sml);
        }
    }
}
void calc(int ct) {
    vector<long long> cn(8, 0), sm(8, 0);
    for (E& e : g[ct]) {
        if (vs[e.v]) continue;
        vector<long long> tcn(8, 0), tsm(8, 0);
        int im = (1 << c[ct]) | (1 << c[e.v]);
        long long ism = e.w % MOD;
        go(e.v, ct, im, ism, tcn, tsm);
        for (int s = 0; s < 8; s++) {
            if (!tcn[s]) continue;
            for (int t = 0; t < 8; t++) {
                if (!cn[t]) continue;
                if ((s | t) == 7) {
                    long long cnt = (tcn[s] * sm[t] % MOD + cn[t] * tsm[s] % MOD) % MOD;
                    res = (res + cnt) % MOD;
                }
            }
        }
        for (int s = 0; s < 8; s++) {
            cn[s] = (cn[s] + tcn[s]) % MOD;
            sm[s] = (sm[s] + tsm[s]) % MOD;
        }
    }
}
void div(int u) {
    int tot = sz(u, -1);
    int ct = -1;
    find(u, -1, tot, ct);
    vs[ct] = true;
    calc(ct);
    for (E& e : g[ct]) {
        if (!vs[e.v]) {
            div(e.v);
        }
    }
}
int main() {
    freopen("blossom.in","r",stdin);
    freopen("blossom.out","w",stdout);
    int n;
    cin >> n;
    for (int i = 1; i <= n; i++) {
        cin >> c[i];
    }
    for (int i = 0; i < n - 1; i++) {
        int u, v, w;
        cin >> u >> v >> w;
        g[u].push_back({v, w});
        g[v].push_back({u, w});
    }
    div(1);
    cout << res % MOD << endl;
    return 0;
}