题意：
long long ans = 0;
for(int i = 1; i <= n; i++)
for(int j = i+1; j <= n; j++)
ans += F(i,j);
F(i,j)表示i点到j点路径上所有的点权和。
若i->j路径上存在2条相邻边边权相同则 F(i,j) = 0
问：ans的值。
int乘法爆掉了我也醉了。。。
思路：
和网上的统计边方法不同，这里是用统计点出现的次数来计算
我们计算每个点i 出现的次数，则答案就是 i的次数*i的点权 => dp[i] * a[i]
而i出现的路径起点和终点有4种
1、i的子孙->i的子孙
2、i的子孙->i
3、i到 （非i的子孙（ 即i的祖先节点，兄弟节点和兄弟节点的子孙
4、i的子孙->非i的子孙
所以先计算1,2的情况 ，用dp1[i]记录
3，4的情况用dp2[i]记录
则答案就是 for(int i = 1; i <= n; i++) ans += a[i] * (dp1[i]+dp2[i]);
siz[u] 表示以u为根的子树中有效的节点数，若 u -> v(col = 1) && v -> k(col=1), 则以k为根的子树都不是有效节点
(其中v是u的儿子，k是v的儿子）
mp[u][col]表示以u为根，有效节点中 用颜色为col的边相连的节点个数
#include <map>
using namespace std;
#define N 300100
struct Edge{
int to, col, nex;
}edge[N<<1];
int head[N], edgenum;
void init(){memset(head, -1, sizeof head); edgenum = 0;}
void add(int u, int v, int col){
Edge E = {v, col, head[u]};
edge[edgenum] = E;
head[u] = edgenum++;
}
typedef long long ll;
template <class T>
inline bool rd(T &ret) {
char c; int sgn;
if(c=getchar(),c==EOF) return 0;
while(c!='-'&&(c<'0'||c>'9')) c=getchar();
sgn=(c=='-')?-1:1;
ret=(c=='-')?0:(c-'0');
while(c=getchar(),c>='0'&&c<='9') ret=ret*10+(c-'0');
ret*=sgn;
return 1;
}
template <class T>
inline void pt(T x) {
if(x>9) pt(x/10);
putchar(x%10+'0');
}
int n, k;
ll dp1[N], dp2[N], a[N];
int siz[N];
map<int, int> mp[N], mp2[N];//mp[u][col]表示u子树下 边颜色=col 的有效的点数
void dfs1(int u, int fa){
siz[u] = 0; dp1[u] = 0;
for(int i = head[u]; ~i; i = edge[i].nex) {
int v = edge[i].to; if(v == fa)continue;
dfs1(v, u);
mp[u][edge[i].col] += siz[v] - mp[v][edge[i].col];
siz[u] += siz[v] - mp[v][edge[i].col];
}
ll dou = 0;
for(int i = head[u]; ~i; i = edge[i].nex) {
int v = edge[i].to; if(v == fa)continue;
dou += (ll)(siz[v] - mp[v][edge[i].col]) * (ll)(siz[u] - mp[u][edge[i].col]);
dp1[u] += siz[v] - mp[v][edge[i].col];
}
dp1[u] += dou >> 1;
siz[u]++;
}
void dfs2(int u, int ok, int col, int fa) {
dp2[u] = (ll)(siz[u] - mp[u][col]) * (ll)ok;
for(int i = head[u]; ~i; i = edge[i].nex) {
int v = edge[i].to; if(v == fa)continue;
if(u != fa && edge[i].col == col)
dfs2(v, siz[u] - mp[u][edge[i].col], edge[i].col, u);
else
dfs2(v, ok + siz[u] - mp[u][edge[i].col], edge[i].col, u);
}
}
void solve(){
init();
for(int i = 1; i <= n; i++)rd(a[i]), mp[i].clear(), mp2[i].clear();
for(int i = 1, u, v, d; i < n; i++) {
rd(u);rd(v);rd(d);
add(u,v,d); add(v,u,d);
}
dfs1(1, 1);
dfs2(1, 0, -1, 1);
}
int main() {
while(rd(n)){
solve();
ll ans = 0;
for(int i = 1; i <= n; i++)
ans += a[i]*(dp1[i]+dp2[i]);
pt( ans );putchar('\n');
}
return 0;
}
/*
4
1 10 100 1000
1 2 1
2 3 1
3 4 1
5
1 10 100 1000 10000
1 2 1
2 3 1
3 4 1
2 5 2
11
1 2 3 4 5 6 7 8 9 111 123
1 2 1
1 3 2
2 4 3
2 5 1
3 6 3
3 7 3
5 8 1
5 9 2
8 10 1
11 8 2
14
1 2 3 4 5 6 7 8 9 111 123 235 66 1000
1 2 1
1 3 2
2 4 3
2 5 1
3 6 3
3 7 3
5 8 1
5 9 2
8 10 1
11 8 2
12 11 2
8 13 1
8 14 2
10
1 1 1 1 1 1 1 1 1 1
1 2 1
1 7 1
1 10 2
2 3 5
2 6 4
3 4 1
3 5 8
7 8 2
7 9 1
14
1 2 5 10 20 30 70 80 100 1000 2000 5000 100000 1000000
1 2 2
1 3 1
1 4 1
2 5 2
2 8 3
3 9 3
3 6 2
4 7 1
4 10 3
3 11 3
6 12 2
13 1 2
14 3 3
*/
时间： 2024-11-05 15:59:37