- 软件开发
- 2019 计蒜之道第一场 C. 商汤AI园区的n个路口（中等）预处理gcd树形dp
商汤AI园区的n个路口（中等）
思路：设dp[u][i]为u的子树中，点u权值为为i的合法方案数，对于u的儿子v，假设w=3，m=6，显然当a[u]=3时，a[v]!=3|6，那么d[u][3]*=(d[v][1]+d[v][2]+d[v][4]+d[v][5])，我们换一种思路，先求出所有的d[v][i]的和res，对于所有的 i，j，如果gcd(i,j)=w，那么我们就用sum[i]+=d[v][j]，然后d[u][i]*=(res-sum[i])，所以我们只要提前预处理所有的gcd(i,j)就好了。
#include<bits/stdc++.h>
#define pi pair<int,int>
#define mk make_pair
#define ll long long
using namespace std
;
const int maxn
=1005,mod
=1e9+7;
vector
<pi
>G
[maxn
],G2
[maxn
];
int n
,m
;
ll ans
,d
[maxn
][maxn
],sum
[maxn
];
void add(ll
& x
,ll y
)
{
x
=(x
+y
)%mod
;
}
ll
dfs(int u
,int fa
)
{
for(int i
=1;i
<=m
;i
++)
d
[u
][i
]=1;
for (auto tmp
: G
[u
]) {
int v
=tmp
.first
;
int w
=tmp
.second
;
if(v
==fa
)
continue;
dfs(v
,u
);
ll res
=0;
for(int i
=1;i
<=m
;i
++)
add(res
,d
[v
][i
]);
for(auto V
: G2
[w
]) {
int i
=V
.first
;
int j
=V
.second
;
add(sum
[i
],d
[v
][j
]);
}
for(int i
=1;i
<=m
;i
++)
d
[u
][i
]=d
[u
][i
]*(res
-sum
[i
]+mod
)%mod
,sum
[i
]=0;
}
}
int main()
{
int u
,v
,w
;
cin
>>n
>>m
;
for(int i
=1;i
<n
;i
++)
{
cin
>>u
>>v
>>w
;
G
[u
].push_back(mk(v
,w
));
G
[v
].push_back(mk(u
,w
));
}
for(int i
=1;i
<=m
;i
++)
for(int j
=1;j
<=m
;j
++)
{
int x
=__gcd(i
,j
);
G2
[x
].push_back(mk(i
,j
));
}
dfs(1,0);
for(int i
=1;i
<=m
;i
++)
add(ans
,d
[1][i
]);
printf("%lld\n",ans
);
}
