test/lowest-common-ancestor.test.cpp

• View this file on GitHub
• Last update: 2021-05-27 04:01:08+09:00
• Problem: https://judge.yosupo.jp/problem/lca

Depends on

• graph/LCA.hpp

Code

#define PROBLEM "https://judge.yosupo.jp/problem/lca"
#include <iostream>
#include <vector>
#include "./graph/LCA.hpp"

int main() {
  std::cin.tie(nullptr);
  std::ios::sync_with_stdio(false);

  int N, Q; std::cin >> N >> Q;
  std::vector<std::vector<int>> g(N);
  for(int i=1; i<N; i++) {
    int p; std::cin >> p;
    g[p].push_back(i);
    g[i].push_back(p);
  }
  LowestCommonAncestor lca(g,0);
  while(Q--) {
    int u, v; std::cin >> u >> v;
    std::cout << lca.query(u, v) << '\n';
  }
  return 0;
}

#line 1 "test/lowest-common-ancestor.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/lca"
#include <iostream>
#include <vector>
#line 2 "graph/LCA.hpp"
#include <cmath>
struct LowestCommonAncestor{
  int n = 0;
  int log2_n = 0;
  std::vector<std::vector<int>> parent;
  std::vector<int> depth;

  LowestCommonAncestor() = default;

  LowestCommonAncestor(const std::vector<std::vector<int>> &g, int root) : n(g.size()), log2_n(std::log2(n) + 1), parent(log2_n, std::vector<int>(n)), depth(n){
    dfs(g, root, -1, 0);
    for(int k=0; k+1 < log2_n;  k++){
      for(int v=0; v<(int)g.size(); v++){
        if(parent[k][v]  < 0) parent[k+1][v] = -1;
        else parent[k+1][v] = parent[k][parent[k][v]];
      }
    }
  }

  void dfs(const std::vector<std::vector<int>> &g, int v, int p, int d) {
    parent[0][v] = p;
    depth[v] = d;
    for(auto &e: g[v]) {
      if(e != p) dfs(g, e, v, d+1);
    }
  }

  int query(int u, int v){
    //uとvの深さが同じになるまで親を辿る
    if(depth[u] > depth[v]) std::swap(u, v);

    for(int k=0; k<log2_n; k++){
      if((depth[v] - depth[u]) >> k & 1) {
        v = parent[k][v];
      }
    }
    if(u == v) return u;
    for(int k=log2_n-1; k>=0; k--) {
      //二分探索でLCAを求める
      if(parent[k][u] != parent[k][v]) {
        u = parent[k][u];
        v = parent[k][v];
      }
    }
    return parent[0][u];
  }
};
#line 5 "test/lowest-common-ancestor.test.cpp"

int main() {
  std::cin.tie(nullptr);
  std::ios::sync_with_stdio(false);

  int N, Q; std::cin >> N >> Q;
  std::vector<std::vector<int>> g(N);
  for(int i=1; i<N; i++) {
    int p; std::cin >> p;
    g[p].push_back(i);
    g[i].push_back(p);
  }
  LowestCommonAncestor lca(g,0);
  while(Q--) {
    int u, v; std::cin >> u >> v;
    std::cout << lca.query(u, v) << '\n';
  }
  return 0;
}

Back to top page