fork download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. using ll = long long;
  5.  
  6. /*
  7. QUESTION IN SIMPLE WORDS
  8. ------------------------
  9. Each node in a tree has an integer value.
  10.  
  11. For a path between nodes u and v:
  12. - collect all values on that path
  13. - find the smallest prime number that does NOT divide all of them
  14.  
  15. Important observation:
  16. A prime divides every number on the path
  17. if and only if that prime divides the GCD of the path.
  18.  
  19. So the real task becomes:
  20. 1) compute GCD of values on the path
  21. 2) find the smallest prime that does not divide that GCD
  22.  
  23. ------------------------------------------------------------
  24. BRUTE FORCE IDEA
  25. ------------------------------------------------------------
  26. For each query (u, v):
  27. - find the full path from u to v
  28. - compute gcd of all values on that path
  29. - check primes 2, 3, 5, 7, ...
  30.  
  31. ------------------------------------------------------------
  32. WHY BRUTE FORCE IS TOO SLOW
  33. ------------------------------------------------------------
  34. If each query walks the tree directly, one query can take O(n).
  35.  
  36. For many queries, that becomes too slow.
  37.  
  38. ------------------------------------------------------------
  39. OPTIMIZED IDEA
  40. ------------------------------------------------------------
  41. Use binary lifting.
  42.  
  43. Binary lifting is normally used to jump upward in a tree quickly.
  44. Here we also store gcd information while jumping upward.
  45.  
  46. Then for a query:
  47. - bring both nodes to same depth
  48. - lift both until they meet
  49. - combine gcd values along the way
  50.  
  51. At the end we get the gcd of the whole path.
  52. Then we return the first prime that does not divide it.
  53. */
  54.  
  55. struct PathPrimeQuery {
  56.     int n, LOG;
  57.     vector<ll> val;
  58.     vector<vector<int>> g;
  59.     vector<int> depth;
  60.  
  61.     // up[k][u] = 2^k-th ancestor of u
  62.     vector<vector<int>> up;
  63.  
  64.     // gup[k][u] = gcd of values along that jump
  65.     vector<vector<ll>> gup;
  66.  
  67.     PathPrimeQuery(int n, const vector<ll>& val) : n(n), val(val) {
  68.         g.assign(n + 1, {});
  69.         depth.assign(n + 1, 0);
  70.  
  71.         LOG = 1;
  72.         while ((1 << LOG) <= n) LOG++;
  73.  
  74.         up.assign(LOG, vector<int>(n + 1, 0));
  75.         gup.assign(LOG, vector<ll>(n + 1, 0));
  76.     }
  77.  
  78.     void addEdge(int u, int v) {
  79.         g[u].push_back(v);
  80.         g[v].push_back(u);
  81.     }
  82.  
  83.     void dfs(int u, int p) {
  84.         up[0][u] = p;
  85.  
  86.         // For one upward jump:
  87.         // if parent does not exist, just use val[u]
  88.         // otherwise gcd(val[u], val[parent])
  89.         if (p == 0) gup[0][u] = val[u];
  90.         else gup[0][u] = __gcd(val[u], val[p]);
  91.  
  92.         for (int k = 1; k < LOG; k++) {
  93.             int mid = up[k - 1][u];
  94.             up[k][u] = up[k - 1][mid];
  95.  
  96.             // Combine two half jumps
  97.             gup[k][u] = __gcd(gup[k - 1][u], gup[k - 1][mid]);
  98.         }
  99.  
  100.         for (int v : g[u]) {
  101.             if (v == p) continue;
  102.             depth[v] = depth[u] + 1;
  103.             dfs(v, u);
  104.         }
  105.     }
  106.  
  107.     void build(int root = 1) {
  108.         dfs(root, 0);
  109.     }
  110.  
  111.     ll pathGcd(int u, int v) {
  112.         ll res = 0;
  113.  
  114.         // Make sure u is deeper
  115.         if (depth[u] < depth[v]) swap(u, v);
  116.  
  117.         // Step 1: lift u up until both nodes are at same depth
  118.         int diff = depth[u] - depth[v];
  119.         for (int k = 0; k < LOG; k++) {
  120.             if (diff & (1 << k)) {
  121.                 res = __gcd(res, gup[k][u]);
  122.                 u = up[k][u];
  123.             }
  124.         }
  125.  
  126.         // If both are same node now, include its value and finish
  127.         if (u == v) {
  128.             res = __gcd(res, val[u]);
  129.             return res;
  130.         }
  131.  
  132.         // Step 2: lift both until they are just below LCA
  133.         for (int k = LOG - 1; k >= 0; k--) {
  134.             if (up[k][u] != up[k][v]) {
  135.                 res = __gcd(res, gup[k][u]);
  136.                 res = __gcd(res, gup[k][v]);
  137.                 u = up[k][u];
  138.                 v = up[k][v];
  139.             }
  140.         }
  141.  
  142.         // Now u and v are children of LCA
  143.         res = __gcd(res, gup[0][u]);
  144.         res = __gcd(res, gup[0][v]);
  145.  
  146.         return res;
  147.     }
  148.  
  149.     bool isPrime(int x) {
  150.         if (x < 2) return false;
  151.         for (int d = 2; 1LL * d * d <= x; d++) {
  152.             if (x % d == 0) return false;
  153.         }
  154.         return true;
  155.     }
  156.  
  157.     int smallestMissingPrime(ll gVal) {
  158.         int p = 2;
  159.         while (true) {
  160.             if (isPrime(p) && gVal % p != 0) return p;
  161.             p++;
  162.         }
  163.     }
  164.  
  165.     int query(int u, int v) {
  166.         ll gVal = pathGcd(u, v);
  167.         return smallestMissingPrime(gVal);
  168.     }
  169. };
  170.  
  171. int main() {
  172.     int n = 5;
  173.     vector<ll> val = {0, 30, 42, 70, 105, 14};
  174.  
  175.     PathPrimeQuery solver(n, val);
  176.     solver.addEdge(1, 2);
  177.     solver.addEdge(1, 3);
  178.     solver.addEdge(2, 4);
  179.     solver.addEdge(2, 5);
  180.  
  181.     solver.build(1);
  182.  
  183.     cout << solver.query(4, 5) << '\n';
  184.     cout << solver.query(3, 4) << '\n';
  185.  
  186.     return 0;
  187. }
Success #stdin #stdout 0.01s 5272KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
2
2
https://ideone.com/m5HagM
language:
C++14 (gcc 8.3)
created:
4 days 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!