首页 > 代码库 > HDU 4303 Hourai Jeweled 树形dp 所有路径点权和 dfs2次

HDU 4303 Hourai Jeweled 树形dp 所有路径点权和 dfs2次

题意:

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


*/


HDU 4303 Hourai Jeweled 树形dp 所有路径点权和 dfs2次