根が頻繁に変わる木で親を求める

根が  r のとき頂点 u の親を  O(\log n) で求めるコード。以下の場合分けによるアルゴリズム

  • $r$ が $u$ の子孫だった場合、$r$ を $u$ にぶつからない限界まで持ち上げる。
  • それ以外の場合、根が頂点 0 だった場合の親をそのまま返す。
// Verified at http://codeforces.com/problemset/problem/916/E
struct HLD {
  using Tree = std::vector<std::vector<int>>;
  std::vector<int> parent, head, vid, inv;

  HLD(const Tree &g) : parent(g.size()), head(g.size()), vid(g.size()), inv(g.size()) {
    int k = 0;
    std::vector<int> size(g.size(), 1);
    std::function<void(int, int)> dfs = [&](int u, int p) {
      for (int v : g[u]) if (v != p) dfs(v, u), size[u] += size[v];
    };
    std::function<void(int, int, int)> dfs2 = [&](int u, int p, int h) {
      parent[u] = p; head[u] = h; vid[u] = k++; inv[vid[u]] = u;
      for (int v : g[u]) if (v != p && size[u] < size[v] * 2) dfs2(v, u, h);
      for (int v : g[u]) if (v != p && size[u] >= size[v] * 2) dfs2(v, u, v);
    };
    dfs(0, -1);
    dfs2(0, -1, 0);
  }

  int find_parent(int u, int r) {
    if (u == r) return -1;
    while (vid[u] < vid[r]) {
      if (head[u] == head[r]) return inv[vid[u] + 1];
      if (parent[head[r]] == u) return head[r];
      r = parent[head[r]];
    }
    return parent[u];
  }
};