fork download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3. #define ll long long int
  4. const int M = 1e9 + 7;
  5.  
  6. int g(int n)
  7. {
  8.   int val = n;
  9.   int cnt = 0;
  10.   int ans = 1;
  11.   while (val % 2 == 0)
  12.   {
  13.     cnt++;
  14.     val /= 2;
  15.   }
  16.   if (cnt % 2 == 1)
  17.     ans *= 2;
  18.   for (int i = 3; i * i <= n; i += 2)
  19.   {
  20.     int cnt = 0;
  21.     while (val % i == 0)
  22.     {
  23.       val /= i;
  24.       cnt++;
  25.     }
  26.     if (cnt % 2 == 1)
  27.       ans *= i;
  28.   }
  29.   if (val != 1)
  30.     ans *= val;
  31.  
  32.   return ans;
  33. }
  34. int dfs(int node, vector<int> &val, vector<int> adj[], unordered_map<int, int> &mp, vector<int> &vis, int sum)
  35. {
  36.   int ans = 0;
  37.   sum += mp[val[node]];
  38.   mp[val[node]]++;
  39.   vis[node] = 1;
  40.   for (auto it : adj[node])
  41.   {
  42.     if (vis[it] == 0)
  43.     {
  44.       ans += dfs(it, val, adj, mp, vis, sum);
  45.     }
  46.   }
  47.   vis[node] = 0;
  48.   mp[val[node]]--;
  49.   if (ans == 0)
  50.   {
  51.     return sum;
  52.   }
  53.   return ans;
  54. }
  55.  
  56. int main()
  57. {
  58.   ios_base::sync_with_stdio(false);
  59.   cin.tie(NULL);
  60.  
  61.   int n;
  62.   cin >> n;
  63.   vector<int> val(n);
  64.   for (int i = 0; i < n; i++)
  65.     cin >> val[i];
  66.   vector<int> adj[n];
  67.   for (int i = 0; i < n - 1; i++)
  68.   {
  69.     int u, v;
  70.     cin >> u >> v;
  71.     adj[u].push_back(v);
  72.     adj[v].push_back(u);
  73.   }
  74.  
  75.   for (int i = 0; i < n; i++)
  76.   {
  77.     int newval = g(val[i]);
  78.     val[i] = newval;
  79.   }
  80.   unordered_map<int, int> mp;
  81.   vector<int> vis(n, 0);
  82.   cout << dfs(0, val, adj, mp, vis, 0);
  83.  
  84.   return 0;
  85. }
Success #stdin #stdout 0.01s 5292KB
comments (?)
stdin
copy 
6 
4 2 4 14 42 1
0 1
1 2
0 3 
3 4 
1 5
stdout
copy 
2
https://ideone.com/5LQO58
language:
C++14 (gcc 8.3)
created:
1 day 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!