hdu_4303_Hourai Jeweled

总共研究了七-八个小时,终于弄出来了,开心呀.一道难题能顶十道水题这句话确实有道理.

树形DP+子树合并,一次DFS即可完成.中等难度,难点在子树合并的细节.

题目大意:  http://acm.hdu.edu.cn/showproblem.php?pid=4303

给一棵树,每个结点有个权值,每个边有颜色,求整棵树中相邻边颜色不同的路径的总权值的合.值得注意的是像题目中 1->2->6 这样一个路径,权值的计算应为 (1+2)+(2+6)+(1+2+6),每一条子路径也应该算上.题目输出数据那里已经解释的很清楚了.

解题思路:

建立模型:

1.点多边少,用邻接数组edge[]存数据,结构体Edge代表从u到v有一个条边颜色是c.

2.数组:

nodevalue: 每个点的初始权值

value:          dp后的总权值

number:     子节点的个数

结果的值为ans;

竖向DP:

1. 对于叶子结点k,value[k]=nodevalue[k],number[k]=1;

2.统计每个结点k的所有子节点个数:  number[k]=sum(number[i])  i为k的所有子节点,并且边i到k的颜色与k到父节点的颜色不同.

3.每个结点k的总权值value[k]为: 所有子节点的总权值合+所有子节点的个数*当前结点的权值.

   value[k]=sum(value[i]+number[i]*nodevalue[k])   i为k的所有子节点,并且边i到k的颜色与k到父节点的颜色不同.

横向合并,并且求值:

1.  每一个结点,对其所有不同颜色的边合并,但这样时间复杂度为n^2,技巧为对一个结点的所有颜色排序,统计与当前颜色不同的颜色的总子节点数sum_n和总权值sum_v后计算.

2. 将每一条边合并后的值相加.即为总值ans.

3. 对于最终结果ans,如果这个结点不是叶子结点,ans要加上当前结点所有子节点的总权值value[k]和当前结点的权值*所有子节点的个数:

number[k]*nodevalue[k];

源代码:

#include <myhead.h>

const int N=300010;
struct Edge{
	static int number;
	int u,v,c;
	void add(int u,int v,int c) {
		this->u=u;
		this->v=v;
		this->c=c;
		Edge::number++;
	}
};
int Edge::number=0;
int n,root;
_I64 ans;
int nodevalue[N],head[N];
_I64 value[N],number[N];
Edge edge[N<<1];

inline void init() {
	Edge::number=root=ans=0;
	memset(head,-1,sizeof(head));
	memset(value,0,sizeof(value));
	memset(number,0,sizeof(number));
	edge[Edge::number].add(0,0,-1);
	nodevalue[root]=0;
	for(int i=1;i<=n;++i) {
		scanf("%d",&nodevalue[i]);
	}
	int u,v,c;
	for(int i=1;i<n;++i) {
		scanf("%d%d%d",&u,&v,&c);
		edge[Edge::number].add(u,v,c);
		edge[Edge::number].add(v,u,c);
	}
}

bool cmp1(const Edge &x,const Edge &y) {
	return x.u<y.u;
}

bool cmp2(const Edge &x,const Edge &y) {
	return x.c<y.c;
}

inline void preEdge() {
	sort(edge,edge+Edge::number,cmp1);
	int start=1;
	head[root]=0;
	head[edge[start].u]=start;
	edge[root].v=edge[start].u;
	for(int i=2;i<Edge::number;++i) {
		if(edge[i].u!=edge[i-1].u) {
			sort(edge+start,edge+i,cmp2);
			start=head[edge[i].u]=i;
		}
	}
	sort(edge+start,edge+Edge::number,cmp2);
}


void dfs(int x,int father,int fathercolor) {
	_I64 fristValue=0,fristNumber=0;
	_I64 lastValue=0,lastNumber=0;
	int lastColor=-1;
	value[x]=nodevalue[x];
	number[x]=1;
	if(head[x]!=-1) {
		for(int i=head[x];edge[i].u==x;++i) {
			if(edge[i].v==father) continue;
			dfs(edge[i].v,x,edge[i].c);
			if(edge[i].c!=fathercolor) {
				number[x]+=number[edge[i].v];
				value[x]+=value[edge[i].v];
				value[x]+=number[edge[i].v]*nodevalue[x];
			}
			if(lastColor!=edge[i].c) {
				lastColor=edge[i].c;
				lastValue=fristValue;
				lastNumber=fristNumber;
			}
			fristValue+=value[edge[i].v];
			fristNumber+=number[edge[i].v];
			ans+=lastValue*number[edge[i].v];
			ans+=lastNumber*value[edge[i].v];
			ans+=lastNumber*number[edge[i].v]*nodevalue[x];
		}
	}
	ans+=value[x]*(1&&father);
	ans+=number[x]*nodevalue[father];
}

int main() {
	while(~scanf("%d",&n)) {
		init();
		preEdge();
		dfs(edge[root].v,root,edge[root].c);
		printf("%I64d\n",ans);
	}
	return 0;
}