比赛 |
2025.3.29 |
评测结果 |
AWWWWWWWTTTTTTTTTTTTTTTTT |
题目名称 |
硝华流焰 |
最终得分 |
4 |
用户昵称 |
LikableP |
运行时间 |
33.909 s |
代码语言 |
C++ |
内存使用 |
5.51 MiB |
提交时间 |
2025-03-29 10:44:26 |
显示代码纯文本
#include <cstdio>
template <typename T>
T max(T x, T y) {
return x > y ? x : y;
}
const int MAXN = 1e5 + 10;
const int MOD = 1e9 + 7;
struct EDGE {
int v, w, next;
} edge[MAXN << 1];
int head[MAXN], edgeNum;
void AddEdge(int u, int v, int w) {
edgeNum++;
edge[edgeNum].v = v;
edge[edgeNum].w = w;
edge[edgeNum].next = head[u];
head[u] = edgeNum;
}
int siz[MAXN], sum;
bool del[MAXN];
void CalcSumAndSize(int root, int fa) {
siz[root] = 1, sum++;
for (int i = head[root]; i; i = edge[i].next) {
int v = edge[i].v;
if (v == fa || del[v]) continue;
CalcSumAndSize(v, root);
siz[root] += siz[v];
}
}
int centroid, depth[MAXN];
void CalcCentroid(int root, int fa) {
int maxx = sum - siz[root];
depth[root] = depth[fa] + 1;
for (int i = head[root]; i; i = edge[i].next) {
int v = edge[i].v;
if (v == fa || del[v]) continue;
CalcCentroid(v, root);
maxx = max(maxx, siz[v]);
}
if (maxx <= sum / 2) {
centroid = root;
}
}
int zero[MAXN], one[MAXN], two[MAXN], zerodepth[MAXN], onedepth[MAXN], twodepth[MAXN], zerocnt, onecnt, twocnt;
int a[MAXN], ans;
void Count(int root, int fa, int c, int nzero, int none, int ntwo) {
c %= MOD;
if (a[root] == 0) {
nzero++;
} else if (a[root] == 1) {
none++;
} else if (a[root] == 2) {
ntwo++;
}
if (nzero && none && ntwo) { // no missing
for (int i = 1; i <= zerocnt; ++i) {
(ans += zero[i] + c) %= MOD;
}
for (int i = 1; i <= onecnt; ++i) {
(ans += one[i] + c) %= MOD;
}
for (int i = 1; i <= twocnt; ++i) {
(ans += two[i] + c) %= MOD;
}
} else if ((!nzero && zerocnt) && none && ntwo) { // miss 0
for (int i = 1; i <= zerocnt; ++i) {
(ans += zero[i] + c) %= MOD;
}
for (int i = 1; i <= onecnt; ++i) {
if (onedepth[i] > zerodepth[zerocnt]) {
(ans += one[i] + c) %= MOD;
}
}
for (int i = 1; i <= twocnt; ++i) {
if (twodepth[i] > zerodepth[zerocnt]) {
(ans += two[i] + c) %= MOD;
}
}
} else if (nzero && (!none && onecnt) && ntwo) { // miss 1
for (int i = 1; i <= zerocnt; ++i) {
if (zerodepth[i] > onedepth[onecnt]) {
(ans += zero[i] + c) %= MOD;
}
}
for (int i = 1; i <= onecnt; ++i) {
(ans += one[i] + c) %= MOD;
}
for (int i = 1; i <= twocnt; ++i) {
if (twodepth[i] > onedepth[onecnt]) {
(ans += two[i] + c) %= MOD;
}
}
} else if (nzero && none && (!ntwo && twocnt)) { // miss 2
for (int i = 1; i <= zerocnt; ++i) {
if (zerodepth[i] > twodepth[twocnt]) {
(ans += zero[i] + c) %= MOD;
}
}
for (int i = 1; i <= onecnt; ++i) {
if (onedepth[i] > twodepth[twocnt]) {
(ans += one[i] + c) %= MOD;
}
}
for (int i = 1; i <= twocnt; ++i) {
(ans += two[i] + c) %= MOD;
}
} else if ((!nzero && zerocnt) && (!none && onecnt) && ntwo) { // miss 0 1
for (int i = 1; i <= zerocnt; ++i) {
if (zerodepth[i] > onedepth[1]) {
(ans += zero[i] + c) %= MOD;
}
}
for (int i = 1; i <= onecnt; ++i) {
if (onedepth[i] > zerodepth[1]) {
(ans += one[i] + c) %= MOD;
}
}
for (int i = 1; i <= twocnt; ++i) {
if (twodepth[i] > max(zerodepth[zerocnt], onedepth[onecnt])) {
(ans += two[i] + c) %= MOD;
}
}
} else if ((!nzero && zerocnt) && none && (!ntwo && twocnt)) { // miss 0 2
for (int i = 1; i <= zerocnt; ++i) {
if (zerodepth[i] > twodepth[1]) {
(ans += zero[i] + c) %= MOD;
}
}
for (int i = 1; i <= onecnt; ++i) {
if (onedepth[i] > max(zerodepth[zerocnt], twodepth[zerocnt])) {
(ans += one[i] + c) %= MOD;
}
}
for (int i = 1; i <= twocnt; ++i) {
if (twodepth[i] > zerodepth[1]) {
(ans += two[i] + c) %= MOD;
}
}
} else if (nzero && (!none && onecnt) && (!ntwo && twocnt)) { // miss 1 2
for (int i = 1; i <= zerocnt; ++i) {
if (zerodepth[i] > max(onedepth[onecnt], twodepth[twocnt])) {
(ans += zero[i] + c) %= MOD;
}
}
for (int i = 1; i <= onecnt; ++i) {
if (onedepth[i] > twodepth[1]) {
(ans += one[i] + c) %= MOD;
}
}
for (int i = 1; i <= twocnt; ++i) {
if (twodepth[i] > onedepth[1]) {
(ans += two[i] + c) %= MOD;
}
}
}
for (int i = head[root]; i; i = edge[i].next) {
int v = edge[i].v, w = edge[i].w;
if (v == fa || del[v]) continue;
Count(v, root, c + w, nzero, none, ntwo);
}
}
void Update(int root, int fa, int c) {
c %= MOD;
if (a[root] == 0) {
zero[++zerocnt] = c;
zerodepth[zerocnt] = depth[root];
} else if (a[root] == 1) {
one[++onecnt] = c;
onedepth[onecnt] = depth[root];
} else {
two[++twocnt] = c;
twodepth[twocnt] = depth[root];
}
for (int i = head[root]; i; i = edge[i].next) {
int v = edge[i].v, w = edge[i].w;
if (v == fa || del[v]) continue;
Update(v, root, c + w);
}
}
void Solve(int root) {
sum = 0;
CalcSumAndSize(root, 0);
CalcCentroid(root, 0);
root = centroid;
del[root] = 1;
zerocnt = onecnt = twocnt = 0;
if (a[root] == 0) {
zero[zerocnt = 1] = 0;
zerodepth[1] = 1;
} else if (a[root] == 1) {
one[onecnt = 1] = 0;
onedepth[1] = 1;
} else {
two[twocnt = 1] = 0;
twodepth[1] = 1;
}
for (int i = head[root]; i; i = edge[i].next) {
int v = edge[i].v, w = edge[i].w;
if (del[v]) continue;
Count(v, root, w, 0, 0, 0);
Update(v, root, w);
}
for (int i = head[root]; i; i = edge[i].next) {
int v = edge[i].v;
if (del[v]) continue;
Solve(v);
}
}
int n;
int main() {
freopen("blossom.in", "r", stdin);
freopen("blossom.out", "w", stdout);
scanf("%d", &n);
for (int i = 1; i <= n; ++i) {
scanf("%d", &a[i]);
}
for (int i = 1; i <= n - 1; ++i) {
int u, v, w;
scanf("%d %d %d", &u, &v, &w);
AddEdge(u, v, w);
AddEdge(v, u, w);
}
Solve(1);
printf("%d\n", ans % MOD);
return 0;
}