商汤AI园区的n个路口（中等） 树形dp

商汤AI园区的n个路口（中等）
题意：有n个点，n-1条边，要求每条边固定一个值为x，与这条边相邻的两个点的权值的gcd不等于x，权值的上限为m，然后问你有多少方案数，使之全成立。
题解：设dp[u][i]代表u节点权值为i的方案数，转移就是u节点的子节点v来转移，如果此时m=4,x=2；那么dp[u][2] = dp[v][1]+dp[v][3];我们dp[v][j]，j不能取2/4,那么我们需要预处理全部的gcd(i,j)=x的对数，剩下的就很简单了。

#include<bits/stdc++.h>
#define ll long long
#define pi pair<int,int>
#define mk make_pair
using namespace std;
const int mod =  1e9+7;
int n,m;
vector<pi>G[1100],sf[1100];
ll ans,dp[1100][1100],sum[1100];
void add(ll &x,ll y)
{
    x += y ;
    if(x>=mod) x-= mod;
    if(x<0) x+= mod;
}
void dfs(int u,int fa){
   for(int i=1;i<=m;i++) dp[u][i] = 1;
   for(auto tmp :G[u])
   {
       int v = tmp.first,w = tmp.second;
       if(v==fa) continue;
       dfs(v,u);
       ll res = 0;
       for(int i=1;i<=m;i++)
        add(res,dp[v][i]);
       for(auto x:sf[w])
       {
           int i = x.first , j = x.second;
           add(sum[i],dp[v][j]);
       }
       for(int i=1;i<=m;i++)
        dp[u][i] = dp[u][i] * ((res-sum[i]+mod)%mod) % mod,sum[i]=0;
   }
}
int main()
{
      scanf("%d%d",&n,&m);
      int u,v,w;
      for(int i=1;i<n;i++)
      {
          scanf("%d%d%d",&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);
          sf[x].push_back(mk(i,j));
      }
      dfs(1,0);
      for(int i=1;i<=m;i++) add(ans,dp[1][i]);
      printf("%lld\n",ans);
}