记录编号 594481 评测结果 AAAAAAAAAAAAAAAAAAAAAAAAA
题目名称 4023.硝华流焰 最终得分 100
用户昵称 ┭┮﹏┭┮ 是否通过 通过
代码语言 C++ 运行时间 1.217 s
提交时间 2024-10-01 11:01:08 内存使用 15.20 MiB
显示代码纯文本
#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define pii pair<int,int>
#define fi first
#define in inline
#define se second
#define mp make_pair
#define pb push_back
const int N = 1e5+10;
const ll inf = 1e18,mod = 1e9+7;

ll read(){
	ll x = 0,f = 1;char c = getchar();
	for(;c < '0' || c > '9';c = getchar())if(c == '-')f = -1;
	for(;c >= '0' && c <= '9';c = getchar())x = (x<<1) + (x<<3) + c-'0';
	return x * f;
}


int n;
int a[N];

vector<pii>e[N];


ll ans = 0,id[3] = {1,2,4};
struct Tree{
	int son[N],top[N],rnk[N],sum[N][3]; 
	int dfn[N],siz[N],fa[N],cnt;
	ll s[8],c[8],dep[N];
	int B(int x,int fa,int bit){return sum[x][bit] - sum[fa][bit] > 0;}
	void dfs1(int x){
		siz[x] = 1,sum[x][a[x]]++;
		for(auto it : e[x]){
			int y = it.fi;
			if(y == fa[x])continue;
			dep[y] = dep[x] + it.se,fa[y] = x;
			sum[y][0] = sum[x][0],sum[y][1] = sum[x][1],sum[y][2] = sum[x][2];
			dfs1(y),siz[x] += siz[y];
			if(!son[x] || siz[y] > siz[son[x]])son[x] = y;
		}
	}
	void dfs2(int x,int t){
		dfn[x] = ++cnt,rnk[cnt] = x,top[x] = t;
		if(son[x])dfs2(son[x],t);
		for(auto it : e[x])
			if(it.fi != son[x] && it.fi != fa[x])dfs2(it.fi,it.fi);
	}
	void build(){dfs1(1),dfs2(1,1);}
	void dfs(int x,bool op){
		for(auto it : e[x])
			if(it.fi != son[x] && it.fi != fa[x])dfs(it.fi,0);
		if(son[x]){
			dfs(son[x],1);
			ll z = 0;
			for(auto it : e[x])
				if(it.fi == son[x])z = it.se;
			for(int i = 0;i < 8;i++)(s[i] += z * c[i]) %= mod;
			ll su[8] = {0},cu[8] = {0};
			for(int i = 0;i < 8;i++)su[i | id[a[x]]] += s[i],cu[i | id[a[x]]] += c[i];
			for(int i = 0;i < 8;i++)s[i] = su[i],c[i] = cu[i];
		}
		c[id[a[x]]]++;
		(ans += s[7]) %= mod;
		for(auto it : e[x]){
			int y = it.fi;
			if(y == son[x] || y == fa[x])continue;
			for(int i = dfn[y];i < dfn[y] + siz[y];i++){
				int now = rnk[i];
				int bit = B(now,x,0) + B(now,x,1) * 2 + B(now,x,2) * 4;
				for(int j = 0;j < 8;j++)
					if((bit | j) == 7)(ans += s[j] + (ll)c[j] * (dep[now] - dep[x])) %= mod;
			}
			for(int i = dfn[y];i < dfn[y] + siz[y];i++){
				int now = rnk[i];
				int bit = B(now,fa[x],0) + B(now,fa[x],1) * 2 + B(now,fa[x],2) * 4;
				(s[bit] += dep[now] - dep[x]) %= mod;
				c[bit]++;
			}
		}
		if(!op)memset(s,0,sizeof s),memset(c,0,sizeof c);
	}
}t;


int main(){
	freopen("blossom.in","r",stdin);
	freopen("blossom.out","w",stdout);
	n = read();
	for(int i = 1;i <= n;i++)a[i] = read();
	for(int i = 1;i < n;i++){
		int x = read(),y = read(),z = read();
		e[x].pb({y,z}),e[y].pb({x,z});
	}
	t.build();
	t.dfs(1,0);
	printf("%lld\n",ans);

	return 0;

}
/*
9
0 0 1 2 0 1 2 1 1
1 2 1
1 3 2
1 4 3
2 7 7
2 8 6
3 5 5
3 6 4
7 9 8
*/