比赛 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;
}