fork download 
  1. #include<bits/stdc++.h>
  2. #include<ext/pb_ds/assoc_container.hpp>
  3. using namespace std;
  4. #define mod (long long int)(1e9+7)
  5. using namespace __gnu_pbds;
  6. typedef long long int ll;
  7. typedef tree<ll, null_type, less<ll>, rb_tree_tag,
  8.              tree_order_statistics_node_update>
  9.     indexed_set;
  10.  
  11. ll n,k;
  12. ll val[500009];
  13. vector<ll>v[500009];
  14. ll dp[100009][101];     // dp[i][j] stores number of ways of getting score j from subtree of node i including node i
  15. ll vis[500009];
  16. ll sub[500009];
  17. ll par[500009];
  18.  
  19. void dfs(ll node)
  20. {
  21.  
  22.   vis[node]=1;
  23.   for(auto child : v[node]){
  24.   if(vis[child]==0){
  25.    par[child]=node;
  26.   dfs(child);
  27.   sub[node]+=sub[child];
  28.   }
  29.   }
  30.  
  31. ll temp[101]={0};        // temporary array to store value of dp[node] array
  32. if(val[node]==1)
  33. temp[1]=1;     // if val[node]==1 we can get score 1 in 1 way without using any child(its subtree) of this node
  34. else
  35. temp[0]=1;   // if val[node]==0 we can get score 0 in 1 way without using ant child(its subtree) of this node
  36.  
  37. for(auto child : v[node])
  38. {
  39.   if(child!=par[node])
  40.   {
  41.      // taversing only upto  min(sub[node],(ll)100) because if size of subtree of node is less than 100(let x)
  42.      // then we can not get score >x from this subtree and if  size of its subtree is greater than 100 we need to calculate upto 100 only
  43.      // because k<=100
  44.  
  45.     for(ll i=min(sub[node],(ll)100);i>=0;i--)                 // traversing from backward to forward to avoid extra counting
  46.     {
  47.        for(ll j=0;j<=min(min(sub[child],(ll)100),i);j++){
  48.        temp[i]=(temp[i]+ ((temp[i-j]*dp[child][j])%mod) )%mod;  // if we can get score i-j without using this child(its subtree) in (temp[i-j] / dp[node][i-j]) ways
  49.        }                                                        // and we can get score j in dp[child][j] ways using this child then we can get score i
  50.     }                                                           // in temp[i-j]*dp[chilc][j] ways using this child
  51.   }
  52. }
  53.  
  54. for(ll i=0;i<=min(sub[node],(ll)100);i++)
  55. dp[node][i]=temp[i];
  56.  
  57.  
  58.  
  59. }
  60.  
  61.  
  62.  
  63. int main()
  64. {
  65.  ios_base::sync_with_stdio(false);
  66.     cin.tie(NULL);
  67.  
  68. ll t=1;
  69. cin>>t;
  70. while(t--)
  71. {
  72.  
  73. ll i;
  74. cin>>n>>k;
  75.  
  76. for(i=1;i<=n;i++)
  77. cin>>val[i];
  78.  
  79. for(i=0;i<n-1;i++)
  80. {
  81.   ll a,b;
  82.    cin>>a>>b;
  83.    v[a].push_back(b);
  84.    v[b].push_back(a);
  85. }
  86. for(i=0;i<=n;i++)
  87. {
  88.   sub[i]=1;
  89.   vis[i]=0;
  90.   for(ll j=0;j<101;j++)
  91.    dp[i][j]=0;
  92. }
  93. par[1]=0;
  94.  
  95. dfs(1);
  96.  
  97. ll ans=0;
  98.  
  99. for(i=1;i<=n;i++)
  100. ans=(ans+dp[i][k])%mod;
  101. cout<<ans<<endl;
  102.  
  103.  
  104. for(i=0;i<=n;i++)
  105. v[i].clear();
  106.  
  107. }
  108. return 0;
  109. }
  110.  
Success #stdin #stdout 0s 15380KB
comments (?)
stdin 
3
3 1
1 0 1 
1 2
1 3
5 2
0 1 0 1 1 
1 2
1 3
1 4
2 5
3 0
1 1 0 
1 2
2 3
stdout 
3
5
1
https://ideone.com/Fplthb
language:
C++14 (gcc 8.3)
created:
4 years ago
visibility:
secret

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!