fork download
copy 
  1. #include<bits/stdc++.h>
  2.  
  3. using namespace std;
  4.  
  5. int n, k;
  6. vector<vector<int> > tree;
  7. vector<int> primes[1000005];
  8. bool isnotprime[1000005];
  9. int a[100005];
  10.  
  11. void pre(){
  12.     // FIXME: It's 1e6 not 1e2
  13.     for(int i = 2; i <= 1e2; ++i){
  14.         if(isnotprime[i]) continue;
  15.         primes[i].push_back(i);
  16.         for(int j = i; j <= 1e2; j += i){
  17.             isnotprime[j] = true;
  18.             primes[j].push_back(i);
  19.         }
  20.     }
  21. }
  22.  
  23. map<int, long long> dp[100005][15]; // (node, k, prime)
  24. long long dp1[100005][15];
  25. long long res;
  26.  
  27. void dfs(int root, int parent){
  28.     for(int child : tree[root]){
  29.         if(child != parent) dfs(child, root);
  30.     }
  31.     for(int p : primes[a[root]]){
  32.         for(int i = 0; i < k; ++i){
  33.             if(i == 0){
  34.                 dp[root][i][p] = 1;
  35.                 if(k == 1) res = 1;
  36.                 dp1[root][i] = 1;
  37.             }
  38.             else {
  39.                 dp[root][i][p] = -1e9;
  40.                 dp1[root][i] = -1e9;
  41.             }
  42.         }
  43.     }
  44.     for(int p : primes[a[root]]){
  45.         for(int remaining = 0; remaining < k; ++remaining){
  46.             // FIXME: Currently, it's iterating every child pairs (ch1, ch2)
  47.             // and calculating the best possible path.
  48.             // Can be done with single iteration. Optimize if TLE.
  49.             for(int ch1 : tree[root]){
  50.                 if(ch1 == parent) continue;
  51.                 for(int ch2 : tree[root]){
  52.                     if(ch2 == parent) continue;
  53.                     if(ch1 == ch2){
  54.                         int rem = ((remaining - 1) % k + k) % k;
  55.                         if(remaining == 0){
  56.                                 dp[root][remaining][p] = max(dp[root][remaining][p], 1 + dp1[ch1][rem]);                                
  57.                         }
  58.                         else{
  59.                             if(dp[ch1][rem][p] > 0){
  60.                                 dp[root][remaining][p] = max(dp[root][remaining][p], 1 + dp[ch1][rem][p]);
  61.                             }
  62.                         }
  63.                         if(dp[root][remaining][p] % k == 0) res = max(res, dp[root][remaining][p]);
  64.                         dp1[root][remaining] = max(dp1[root][remaining], dp[root][remaining][p]);
  65.                     }
  66.                     else{
  67.                         int  lefttree_rem = ((remaining - 1) % k + k) % k;
  68.                         int righttree_rem = (k - (((remaining + 2) % k + k) % k)) % k;
  69.                         long long c = -1e9;
  70.                         int x, y;
  71.                         if(remaining == 0 || remaining == k - 1){
  72.                             if(remaining == 0 && remaining == k - 1){
  73.                                 x = dp1[ch1][lefttree_rem];
  74.                                 y = dp1[ch2][righttree_rem];
  75.                                 if(x > 0 && y > 0){
  76.                                     c = 1 + x + y;
  77.                                 }
  78.                             }
  79.                             else if(remaining == 0 && remaining != k - 1){
  80.                                 x = dp1[ch1][lefttree_rem];
  81.                                 y = dp[ch2][righttree_rem][p];
  82.                                 if(x > 0 && y > 0){
  83.                                     c = 1 + x + y;
  84.                                 }
  85.                             }
  86.                             else if(remaining == k - 1 && remaining != 0){
  87.                                 x = dp[ch1][lefttree_rem][p];
  88.                                 y = dp1[ch2][righttree_rem];
  89.                                 if(x > 0 && y > 0){
  90.                                     c = 1 + x + y;
  91.                                 }
  92.                             }
  93.                         }
  94.                         else{
  95.                             x = dp[ch1][lefttree_rem][p];
  96.                             y = dp[ch2][righttree_rem][p];
  97.                             if(x > 0 && y > 0){
  98.                                 c = 1 + x + y;
  99.                             }
  100.                         }
  101.                         if(c % k == 0) res = max(res, c);
  102.                     }
  103.                 }
  104.             }
  105.         }
  106.     }
  107. }
  108.  
  109. int dfs1(int cur, int par){
  110.     vector<long long> p;
  111.     for(int child : tree[cur]){
  112.         if(child != par){
  113.             int ret = dfs1(child, cur);
  114.             p.push_back(ret);
  115.         }
  116.     }
  117.     if(p.size() == 1){
  118.         sort(p.begin(), p.end());
  119.         if(a[cur] > 1) res = max(res, 1 + p[p.size() - 1]);
  120.     }
  121.     else if(p.size() >= 2){
  122.         sort(p.begin(), p.end());
  123.         if(a[cur] > 1) res = max(res, 1 + p[p.size() - 1] + p[p.size() - 2]);
  124.     }
  125.     if(a[cur] > 1){
  126.         if(p.size()) p[p.size() - 1] += 1;
  127.         else p.push_back(1);
  128.         res = max(res, p.back());
  129.     }
  130.     else{
  131.         p.clear();
  132.         p.push_back(0);
  133.     }
  134.     return p[p.size() - 1];
  135. }
  136.  
  137. int main(){
  138.     pre();
  139.     cin >> n >> k;
  140.     tree.resize(n + 5);
  141.     int x, y;
  142.     for(int i = 1; i < n; ++i){
  143.         scanf("%d%d", &x, &y);
  144.         tree[x].push_back(y);
  145.         tree[y].push_back(x);
  146.     }
  147.     for(int i = 1; i <= n; ++i){
  148.         scanf("%d", &a[i]);
  149.     }
  150.     if(k == 1){
  151.         dfs1(1, -1);
  152.         cout << res << endl;
  153.         return 0;
  154.     }
  155.     dfs(1, -1);
  156.     cout << res << endl;
  157.     return 0;
  158. }
Success #stdin #stdout 0.05s 96756KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
0
https://ideone.com/wv8Dj1
language:
C++14 (gcc 8.3)
created:
6 years ago
visibility:
public

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!