问题 B: FZB（树形DP+边记忆化）

问题 B: FZB

时间限制: 1 Sec 内存限制: 128 MB

题目描述

Hja有一棵N个点的树，树上每个点有点权，每条边有颜色。
一条路径的权值是这条路径上所有点的点权和，一条合法的路径需要满足该路径上任意相邻的两条边颜色都不相同。
问这棵树上所有合法路径的权值和是多少。

输入

第一行一个数N。
接下来一行N个数代表每个点的权值。
接下来N−1行每行三个整数s,e,c，代表s到e之间有一条颜色为c的边。

输出

一行一个整数代表答案。
样例输入
6
6 2 3 7 1 4
1 2 1
1 3 2
1 4 3
2 5 1
2 6 2
样例输出
134

提示

对于30%的数据，1≤N≤1000。
对于另外20%的数据，数据随机。
对于另外20%的数据，是一条链。
对于100%的数据，\(1≤N≤3×10^5,1≤c≤10^9\)。

做法：

树形\(DP\)+边记忆化。
每次以i为端点\(dfs\)，同时记忆化从子树中的一条边开始的合法路径的条数及总权值。
这样每条边（视为有向边）最多访问一次，总边数为\(2*n-2\)，所以复杂度看似为\(O(n)\)？
最后，因为从每个点开始都\(dfs\)了一遍，所以结果要除2。

详细过程看代码：

//#pragma GCC optimize(2)
//#pragma GCC optimize(3, "Ofast", "inline")
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int,int> pii;
const int N=1e6+5;
const ll mod=1e9+7;
const double eps=1e-5;
//const double pi=acos(-1);

#define ls p<<1
#define rs p<<1|1
int tot,nex[N],head[N],ver[N],col[N];
ll a[N];
void add(int u,int v,int w)
{
    nex[++tot]=head[u];head[u]=tot;ver[tot]=v;col[tot]=w;
    nex[++tot]=head[v];head[v]=tot;ver[tot]=u;col[tot]=w;
}
struct node
{
    ll sum,cnt;
    node operator+(const node &t)const
    {
        return node{sum+t.sum,cnt+t.cnt};
    }
}vis[N];
node dfs(int u,int fa,ll sum,int e)
{
    node now={sum,1};
    for(int i=head[u];i;i=nex[i])
    {
        int v=ver[i];
        if(v==fa||col[i]==e) continue;
        if(vis[i].cnt) now=node{now.sum+vis[i].cnt*sum+vis[i].sum,now.cnt+vis[i].cnt};
        else
        {
            node tmp=dfs(v,u,sum+a[v],col[i]);
            now=now+tmp;
            vis[i]=node{tmp.sum-sum*tmp.cnt,tmp.cnt};
        }
    }
    return now;
}
int main() {
#ifndef ONLINE_JUDGE
    freopen("in.txt", "r", stdin);
#endif
    ios::sync_with_stdio(false);
    cin.tie(0);
    int n;
    cin>>n;ll ans=0;
    for(int i=1;i<=n;i++)
        cin>>a[i],ans-=a[i];
    for(int i=1;i<n;i++)
    {
        int u,v,w;
        cin>>u>>v>>w;
        add(u,v,w);
    }
    for(int i=1;i<=n;i++)
        ans+=dfs(i,0,a[i],0).sum;
    cout<<ans/2<<endl;
    return 0;
}