peltorator

LCA O(n + q) Farach-Colton and Bender

Apr 6th, 2021 (edited)
193
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #include <bits/stdc++.h>
  2.  
  3. using namespace std;
  4.  
  5. template<typename T>
  6. struct RMQPM1 {
  7.     vector<T> arr;
  8.     vector<vector<T>> sparse;
  9.     vector<vector<vector<int>>> precalc;
  10.     vector<int> blockType;
  11.     vector<int> logs;
  12.     int n;
  13.     int k;
  14.     int blocks;
  15.     int lg;
  16.     T INF;
  17.     function<int(T)> getKey;
  18.  
  19.     RMQPM1() {};
  20.  
  21.     RMQPM1(const vector<T>& arr, const T& infty, function<int(T)> getKey) : arr(arr), getKey(getKey) {
  22.         n = arr.size();
  23.         k = ceil(log2(n + 1) / 2);
  24.         blocks = (n + k - 1) / k;
  25.         lg = ceil(log2(blocks + 1));
  26.         sparse.assign(lg, vector<T>(blocks));
  27.         precalc.assign((1 << (k - 1)), vector<vector<int>>(k, vector<int>(k + 1, 0)));
  28.         blockType.assign(blocks, 0);
  29.         logs.assign(blocks, 0);
  30.         INF = infty;
  31.        
  32.         buildLogs();
  33.         buildSparse();
  34.         buildPrecalc();
  35.     }
  36.  
  37.     void buildLogs() {
  38.         logs[1] = 0;
  39.         for (int i = 2; i < blocks; i++) {
  40.             logs[i] = logs[i >> 1] + 1;
  41.         }
  42.     }
  43.  
  44.     void buildSparse() {
  45.         for (int block = 0; block < blocks; block++) {
  46.             int l = block * k;
  47.             int r = min(l + k, n);
  48.             T blockMin = arr[l];
  49.             for (int i = l + 1; i < r; i++) {
  50.                 blockMin = min(blockMin, arr[i]);
  51.                 if (getKey(arr[i]) == getKey(arr[i - 1]) + 1) {
  52.                     blockType[block] |= (1 << (i - l - 1));
  53.                 }
  54.             }
  55.             sparse[0][block] = blockMin;
  56.         }
  57.  
  58.         for (int bit = 1; bit < lg; bit++) {
  59.             for (int i = 0; i + (1 << bit) <= blocks; i++) {
  60.                 sparse[bit][i] = min(sparse[bit - 1][i], sparse[bit - 1][i + (1 << (bit - 1))]);
  61.             }
  62.         }
  63.     }
  64.  
  65.     void buildPrecalc() {
  66.         for (int mask = 0; mask < (1 << (k - 1)); mask++) {
  67.             for (int l = 0; l < k; l++) {
  68.                 precalc[mask][l][l] = -1;
  69.                 int curmin = numeric_limits<int>::max();
  70.                 int curval = 0;
  71.                 for (int r = l; r < k; r++) {
  72.                     if (curval < curmin) {
  73.                         precalc[mask][l][r + 1] = r;
  74.                         curmin = curval;
  75.                     } else {
  76.                         precalc[mask][l][r + 1] = precalc[mask][l][r];
  77.                     }
  78.                     if ((mask >> r) & 1) {
  79.                         curval++;
  80.                     } else {
  81.                         curval--;
  82.                     }
  83.                 }
  84.             }
  85.         }
  86.     }
  87.    
  88.     T findShort(int left, int right) { // left and right are in one block
  89.         if (left == right) {
  90.             return INF;
  91.         }
  92.         int block = left / k;
  93.         T ans = arr[block * k + precalc[blockType[block]][left - block * k][right - block * k]];
  94.         return ans;
  95.     }
  96.  
  97.     T findLong(int lblock, int rblock) {
  98.         if (lblock == rblock) {
  99.             return INF;
  100.         }
  101.         int bit = logs[rblock - lblock];
  102.         T ans = min(sparse[bit][lblock], sparse[bit][rblock - (1 << bit)]);
  103.         return ans;
  104.     }
  105.  
  106.     T findMin(int left, int right) { // [left; right)
  107.         int lblock = left / k;
  108.         int rblock = right / k;
  109.         if (lblock == rblock) {
  110.             return findShort(left, right);
  111.         } else {
  112.             return min(findLong(lblock + 1, rblock), min(findShort(left, (lblock + 1) * k), findShort(rblock * k, right)));
  113.         }
  114.     }
  115. };
  116.  
  117. struct FCBLCA {
  118.     vector<vector<int>> graph;
  119.     vector<pair<int, int>> eulertour;
  120.     vector<int> timein;
  121.     int n;
  122.     RMQPM1<pair<int, int>> rmq;
  123.  
  124.     FCBLCA(const vector<vector<int>>& graph, int root = 0) : graph(graph) {
  125.         n = graph.size();
  126.         timein.assign(n, 0);
  127.         dfs(root);
  128.         rmq = RMQPM1<pair<int, int>>(eulertour, {1e9 + 7, 0}, [](pair<int, int> pr) { return pr.first; });
  129.     }
  130.  
  131.     void dfs(int v, int depth = 0, int anc = -1) {
  132.         timein[v] = eulertour.size();
  133.         eulertour.emplace_back(depth, v);
  134.         for (int u : graph[v]) {
  135.             if (u != anc) {
  136.                 dfs(u, depth + 1, v);
  137.                 eulertour.emplace_back(depth, v);
  138.             }
  139.         }
  140.     }
  141.  
  142.     int findLCA(int u, int v) {
  143.         if (timein[u] > timein[v]) {
  144.             swap(u, v);
  145.         }
  146.         return rmq.findMin(timein[u], timein[v] + 1).second;
  147.     }
  148. };
  149.  
  150. int main() {
  151.     int n;
  152.     cin >> n;
  153.     vector<vector<int>> graph(n);
  154.     for (int i = 1; i < n; i++) {
  155.         int u, v;
  156.         cin >> u >> v;
  157.         u--;
  158.         v--; // 0-indexing
  159.         graph[u].push_back(v);
  160.         graph[v].push_back(u);
  161.     }
  162.     FCBLCA fcblca(graph);
  163.     int q;
  164.     cin >> q;
  165.     for (int i = 0; i < q; i++) {
  166.         int u, v;
  167.         cin >> u >> v;
  168.         u--;
  169.         v--; // 0-indexing
  170.         cout << fcblca.findLCA(u, v) + 1 << '\n';
  171.     }
  172.     return 0;
  173. }
  174.  
RAW Paste Data