题目大意:给一棵只有0边权的树和1边权的树,问有多少个点对(x,y),使得从x到y的简单路径满足:一旦经过1边权的边,就不再经过0边权的边。
解法:树形dp或树分治
这题一看就是可以树形dp搞或树分治搞,树分治太麻烦了,又不好写,一般都用树形dp取代,除非万不得已。
考虑树形dp的做法:显然还是先对问题进行分解,把路径分成经过根节点的路径,dfs处理所有点为根节点的情况得到答案。考虑每一个子问题:当处理到点i时,答案的贡献从何而来?以i点为根节点,其他点到根节点的路径中,对答案可能有贡献的路径有4种,分别为(00),(01),(10),(11),分别表示一路都是0的边的路径,先连续0后连续1的路径,先连续1后连续0的路径,都是1的路径,中间夹着0或者中间夹着1的都是没用了的(想一想为什么),用dp[i][j]记录i点为根节点,到i点的路径为j类型的点的个数,这四类边可以直接用二进制表示。
如何统计答案呢? 观察到几种类型的点是可以组成符合答案的点对的:(00) -> (01),(00) -> (00),(00) -> (11),(01)-> (00)(这里没有重复,因为是有序的,这和上面那个方向不同),(10) -> (11),(11) -> (11),(11 -> 10),(11) -> (00);
注意这里的点对是有序的,(x,y) 和 (y,x)是不同的点对,注意统计答案的时候小心这点。
(不要忘了统计子节点到根节点的点对树)
第二个关键点就是边权转点权了(也就是如何统计路径类型点的个数):可以把边权转为点权,每个点附上的权是它的父节点指向它的那条边的权,然后再分为那四种类型,注意当一个子树的问题计算完的时候,它的父节点处理子树的问题是,路径性质是有改变的,因为多了一条指向子树的边,万幸的是这里树形dp还是可以处理这种情况的。想一想添加一条边后路径的性质会变成什么(哪些会变成哪些,哪些会保持原样,哪些会作废),然后更改一下信息,使得它可以被当前结点所用。
转移和点的处理每个人的写法可能都不一样,基本思想还是把问题分解,依次处理每个点为根节点的情况。
(这题有更简单的做法,用并查集维护连通块)
(这题不难想,但处理细节繁多(对我来说。。))
#include<bits/stdc++.h>
using namespace std;
#define pii pair<int,int>
const int maxn = 2e5 + 10;
int n;
vector<pii> g[maxn];
long long dp[maxn][4];
long long res = 0;
void dfs(int s,int fa,int pre) {
dp[s][0] = dp[s][1] = dp[s][2] = dp[s][3] = 0;
for(int i = 0; i < g[s].size(); i++) {
int v = g[s][i].first;
int w = g[s][i].second;
if(v == fa) continue;
dfs(v,s,w);
res += dp[v][0] * dp[s][0] * 2;
res += dp[v][0] * dp[s][3];
res += dp[v][0] * dp[s][1];
res += dp[v][1] * dp[s][0];
res += dp[v][2] * dp[s][3];
res += dp[v][3] * dp[s][0];
res += dp[v][3] * dp[s][2];
res += dp[v][3] * dp[s][3] * 2;
if(w == 0) {
dp[s][0] += dp[v][0];
dp[s][1] += dp[v][3] + dp[v][1];
}
else {
dp[s][3] += dp[v][3];
dp[s][2] += dp[v][0] + dp[v][2];
}
}
res += dp[s][0] * 2 + dp[s][1] + dp[s][2] + dp[s][3] * 2; //统计子结点到根结点的答案
if(pre == 1) {
dp[s][3]++; //把s结点加进去
dp[s][2] += dp[s][0];
dp[s][1] = 0;
dp[s][0] = 0;
}
if(pre == 0) {
dp[s][1] += dp[s][3];
dp[s][2] = dp[s][3] = 0;
dp[s][0]++;
}
}
int main() {
scanf("%d",&n);
for(int i = 1; i < n; i++) {
int x,y,v;
scanf("%d%d%d",&x,&y,&v);
g[x].push_back(pii(y,v));
g[y].push_back(pii(x,v));
}
res = 0;
dfs(1,-1,-1);
printf("%lld\n",res);
return 0;
}