题面:https://ac.nowcoder.com/acm/contest/33190/D
分析:
题意相当于给你一棵树,树上的节点有两种颜色,求联通子树中度数为1的点颜色相同的方案数
对于这一道题,我们可以想到用树形DP来求解
我们设
d
p
i
,
c
o
l
dp_{i,col}
dpi,col为以
i
i
i为根的子树内
x
x
x被选择,除
x
x
x外叶节点为
c
o
l
col
col的方案数
可以看出当颜色为
c
o
l
col
col时
d
p
x
,
c
o
l
=
∏
s
ϵ
s
o
n
x
(
d
p
s
o
n
x
,
c
o
l
+
1
)
dp_{x,col}=\prod_{s\epsilon son_x}(dp_{son_{x,col}}+1)
dpx,col=sϵsonx∏(dpsonx,col+1)
对
d
p
s
o
n
x
,
c
o
l
+
1
dp_{son_{x,col}}+1
dpsonx,col+1的解释:因为对于每一个方案可以取
0
至
d
p
s
o
n
x
,
c
o
l
0至dp_{son_{x,col}}
0至dpsonx,col,所以共可取
d
p
s
o
n
x
+
1
dp_{son_x}+1
dpsonx+1 种
代码:
#include<bits/stdc++.h>
using namespace std;
const int MXN=3*1e5+7;
const long long mod=1e9+7;
int n;
vector<int> g[MXN];
long long dp[MXN][3];
long long res=0;
char s[MXN];
inline int read(){
int x=0,f=1;
char ch=getchar();
while(ch<'0'||ch>'9'){
ch=getchar();
}
while(ch>='0'&&ch<='9'){
x=(x<<3)+(x<<1)+(ch^48);
ch=getchar();
}
return x;
}
void dfs(int u,int fa){
dp[u][1]=dp[u][0]=1;
int i,v;
long long s1=0,s0=0;
for(i=0;i<g[u].size();i++){
v=g[u][i];
if(v==fa) continue;
dfs(v,u);
dp[u][1]=(1ll*dp[u][1]*(dp[v][1]+1))%mod;
dp[u][0]=(1ll*dp[u][0]*(dp[v][0]+1))%mod;
s1=(s1+dp[v][1])%mod;
s0=(s0+dp[v][0])%mod;
}
if(s[u]=='1'){
dp[u][0]--;
res=(res+dp[u][1])%mod;
res=(res-s0+dp[u][0])%mod;
} else {
dp[u][1]--;
res=(res+dp[u][0])%mod;
res=(res-s1+dp[u][1])%mod;
}
}
int main(){
int i,x,y;
n=read();
scanf("%s",s+1);
for(i=1;i<n;i++){
x=read(); y=read();
g[x].push_back(y);
g[y].push_back(x);
}
dfs(1,-1);
printf("%lld\n",(res%mod+mod)%mod);
}