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;
}