本题已收录至2019/9/15 本周总结
题目
【问题描述】
Hja有一棵(n)个点的树,树上每个点有点权,每条边有颜色.一条路径的权值是这条路径上所有点的点权和,一条合法的路径需要满足该路径上任意相邻的两条边颜色都不相同.问这棵树上所有合法路径的权值和是多少.
【输入格式】
第一行一个数(n).
接下来一行(n)个数代表每个点的权值.
接下来(n-1)行每行三个整数(u v c)，代表(u)到(v)之间有一条颜色为(c)的边。
【输出格式】
一行一个整数代表答案.
【样例输入】
6
6 2 3 7 1 4
1 2 1
1 3 2
1 4 3
2 5 1
2 6 2
【样例输出】
134
【数据范围与规定】
(1le n le 3 imes 10^5,1le c le 10^9)
保证答案小于(2^{63}-1).
思路1 (O(n^2))暴力
从每个点DFS一次,计算出以该点为端点的所有合法路径的权值和.每条链被计算了两次,因此答案要(/2).
时间复杂度(O(n^2))
#include<bits/stdc++.h>
#define LL long long
const int SIZE=300005;
int n,head[SIZE],nex[SIZE*2],to[SIZE*2],edge[SIZE*2],Tot;
LL weight[SIZE],Ans;
void Link(int u,int v,int e)
{
nex[++Tot]=head[u];head[u]=Tot;to[Tot]=v;edge[Tot]=e;
nex[++Tot]=head[v];head[v]=Tot;to[Tot]=u;edge[Tot]=e;
}
void DFS(int u,int F,LL sum,int Las)
{
Ans+=sum;
for(int i=head[u];i;i=nex[i])
{
int v=to[i];
if(v==F||edge[i]==Las)continue;//不能出现连续两条颜色相同的边
DFS(v,u,sum+weight[v],edge[i]);
}
}
int main()
{
int u,v,e;
scanf("%d",&n);
for(int i=1;i<=n;i++)
{
scanf("%lld",&weight[i]);
Ans-=weight[i];
}
for(int i=1;i<n;i++)
{
scanf("%d%d%d",&u,&v,&e);
Link(u,v,e);
}
for(int i=1;i<=n;i++)DFS(i,0,weight[i],0);
printf("%lld",Ans/2);
return 0;
}
思路2 (O(n^2))记忆化搜索
容易发现,思路1的搜索有大量的重复状态.
举个例子,
现在我们从涂黑的点开始搜索,在粉框里搜出来了4种状态,列举在了右边.
需要注意的是,有一种状态(粉框左下角的连续两条蓝边)是不合法的,我们没有保留这种状态.
现在我们换第二个点进行DFS,就拿上图中涂黑的点来举例吧.
从这个点开始搜索,依然会进入粉框,在粉框中依然会搜出这4种状态,列举在了右边.
容易发现,这4种状态和上一个点的状态基本一样,只有绿框中的不一样.
换句话说,这4种状态在粉框中的部分是一样的.
这启示我们把粉框里搜出来的状态存起来,下次直接调用,不再在粉框内部进行DFS.
也就是这样:
当我们在搜涂黑的点时,不用进入粉框搜索了,直接取出粉框的状态,加入答案.
容易发现,状态应该存储在(有向)边上,而且这样的记忆化可以推广到所有边.
于是我们就得到了记忆化搜索的代码.
#include<bits/stdc++.h>
#define LL long long
const int SIZE=600005;
int n,head[SIZE],nex[SIZE],to[SIZE],edge[SIZE],Tot;
LL weight[SIZE],Ans;
void Link(int u,int v,int e)
{
nex[++Tot]=head[u];head[u]=Tot;to[Tot]=v;edge[Tot]=e;
nex[++Tot]=head[v];head[v]=Tot;to[Tot]=u;edge[Tot]=e;
}
struct Re
{
LL Res;
int w;
Re operator +(const Re &o)const
{
return (Re){Res+o.Res,w+o.w};
}
}x[SIZE];
Re DFS(int u,int F,LL sum,int Las)
{
Re R=(Re){sum,1};
for(int i=head[u];i;i=nex[i])
{
int v=to[i];
if(v==F||edge[i]==Las)continue;
if(x[i].w)R=(Re){R.Res+x[i].w*sum+x[i].Res,R.w+x[i].w};
else
{
Re Tem=DFS(v,u,sum+weight[v],edge[i]);
R=R+Tem;
x[i]=(Re){Tem.Res-sum*Tem.w,Tem.w};
}
}
return R;
}
int main()
{
int u,v,e;
scanf("%d",&n);
for(int i=1;i<=n;i++)
{
scanf("%lld",&weight[i]);
Ans-=weight[i];
}
for(int i=1;i<n;i++)
{
scanf("%d%d%d",&u,&v,&e);
Link(u,v,e);
}
for(int i=1;i<=n;i++)
Ans+=DFS(i,0,weight[i],0).Res;
printf("%lld",Ans/2);
return 0;
}
容易证明:
一条链的情况时间复杂度(O(n)).
随机数据时间复杂度近似(O(n)).
菊花图时间复杂度(O(n^2))
证明略.
思路3 记忆化搜索的优化
刚刚睡觉的时候想到可以对点也记忆化,但是这样要用map,时间复杂度(O(nlogn)).
#include<bits/stdc++.h>
#define LL long long
const int SIZE=600005;
int n,head[SIZE],nex[SIZE],to[SIZE],edge[SIZE],Tot;
LL weight[SIZE],Ans;
void Link(int u,int v,int e)
{
nex[++Tot]=head[u];head[u]=Tot;to[Tot]=v;edge[Tot]=e;
nex[++Tot]=head[v];head[v]=Tot;to[Tot]=u;edge[Tot]=e;
}
struct Re
{
LL Res;int w;
Re operator +(const Re &o)const{return (Re){Res+o.Res,w+o.w};}
Re operator -(const Re &o)const{return (Re){Res-o.Res,w-o.w};}
}x[SIZE],mk[SIZE];
std::map<std::pair<int,int>,Re>mp;
bool Finished[SIZE];
Re DFS(int u,int F,LL sum,int Las)
{
Re R=(Re){sum,1};
if(Finished[u])
{
Re Tem=mk[u]-mp[std::make_pair(u,Las)];
R=(Re){R.Res+Tem.w*sum+Tem.Res,R.w+Tem.w};
return R;
}
for(int i=head[u];i;i=nex[i])
{
int v=to[i];
if(v==F){if(x[i].w)Finished[u]=1;continue;}
if(x[i].w)
{
if(edge[i]!=Las)R=(Re){R.Res+x[i].w*sum+x[i].Res,R.w+x[i].w};
continue;
}
Re Tem=DFS(v,u,sum+weight[v],edge[i]);
if(edge[i]!=Las)R=R+Tem;
x[i]=(Re){Tem.Res-sum*Tem.w,Tem.w};
mp[std::make_pair(u,edge[i])]=mp[std::make_pair(u,edge[i])]+x[i];
mk[u]=mk[u]+x[i];
}
return R;
}
int main()
{
int u,v,e;
scanf("%d",&n);
for(int i=1;i<=n;++i){scanf("%lld",&weight[i]);Ans-=weight[i];}
for(int i=1;i<n;++i){scanf("%d%d%d",&u,&v,&e);Link(u,v,e);}
for(int i=1;i<=n;++i)Ans+=DFS(i,0,weight[i],0).Res;
printf("%lld",Ans/2);
return 0;
}