Cod sursa(job #2976655)

Utilizator gabrielinelusGabriel-Robert Inelus gabrielinelus Data 9 februarie 2023 20:28:28
Problema Lowest Common Ancestor Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.22 kb
#include <cstdio>
#include <vector>
#include <memory>
#include <algorithm>

using namespace std;

class LCA{
private:
  vector<vector<int>> tree;
  vector<vector<int>> DP;
  vector<int> depth;
  int treeSize;
  int logpow2;
  
  // DP[i][j] = 2^ith ancestor of j
  // DP[k][i] = 2^k th parent of i = 2^(k-1)th parent of 2^(k-1)th parent of i

  void DFS(int k, int treeDepth) {
    depth[k] = treeDepth;
    for (auto v: tree[k])
      DFS(v, treeDepth + 1);
  }
  
public:
  LCA(vector<int> parent) {
    treeSize = parent.size();
    logpow2 = 0;
    int pow2 = 1;

    while (pow2 < treeSize) {
      pow2 <<=1;
      ++logpow2;
    }
    
    DP.resize(logpow2 + 1, vector<int>(treeSize));
    for (int i = 0; i < treeSize; ++i) 
      DP[0][i] = parent[i]; // 2^0th parent of i is it's parent, root's parent is itself.
    
    
    for (int k = 1; k <= logpow2; ++k)
      for (int i = 0; i < treeSize; ++i)
	DP[k][i] = DP[k - 1][DP[k - 1][i]];
    
    
    tree.resize(treeSize);
    depth.resize(treeSize);

    for (int i = 1; i < treeSize; ++i)
      tree[parent[i]].emplace_back(i);
    
    DFS(0, 0);
  }
  int queryAncestor(int a, int b) {
    int depthA = depth[a];
    int depthB = depth[b];

    if (depthA < depthB) {
      swap(a, b);
      swap(depthA, depthB);
    }

    if (depthA != depthB) {
      int delta = depthA - depthB;
      for (int i = logpow2; i >= 0; --i)
	if (delta & (1<<i)) {
	  depthA -= (1<<i);
	  a = DP[i][a];
	}
    }

    if (a == b)
      return a;

    for (int i = logpow2; i >= 0; --i) {
      if (DP[i][a] != DP[i][b]) {
	a = DP[i][a];
	b = DP[i][b];
      }
    }
    return DP[0][a];
  }
};

class Solver{
private:
public:
  Solver() {
    freopen("lca.in", "r", stdin);
    freopen("lca.out", "w", stdout);
  }
  ~Solver() {
    fclose(stdin);
    fclose(stdout);
  }
  void execute() {
    int N, Q;
    scanf("%d%d", &N, &Q);
    vector<int> parent(N);
    for (int i = 1; i < N; ++i) {
      scanf("%d", &parent[i]);
      --parent[i];
    }
    LCA lca(parent);

    int a, b;
    for (int i = 0; i < Q; ++i) {
      scanf("%d%d", &a, &b);
      --a; -- b;
      printf("%d\n", lca.queryAncestor(a, b) + 1);
    }
  }
};

int main() {
  unique_ptr<Solver> s = make_unique<Solver>();
  s->execute();
  return 0;
}