木に対する一般的なクエリ
(tree/tree-query.hpp)

• View this file on GitHub
• Last update: 2024-05-03 23:21:26+09:00
• Include: #include "tree/tree-query.hpp"

木に対する一般的なクエリ

概要

木を触る時によく使うクエリを詰め込んだデータ構造。ダブリングをベースに実装している。

余談

nxt(s, t)を$\langle$前計算$\mathrm{O}(N)$,クエリ$\mathrm{O}(1)$$\rangle$で出来るか考えたけど良くわからなかった。$\langle\mathrm{O}(N),\mathrm{O}(1)\rangle$RMQとかで出来るんだろうか…

ちなみにsparse tableを使えば$\langle\mathrm{O}(N\log N),\mathrm{O}(1)\rangle$は達成可能だと思う。自分の実装は$\langle\mathrm{O}(N\log N),\mathrm{O}(\log N)\rangle$だが空間の定数倍がsparse tableより2倍くらいよい(多分)のでMLEに強いはず…

追記：よくよく考え直すとHLDで簡単に$\langle\mathrm{O}(N),\mathrm{O}(\log N)\rangle$が書けないか？(おいおい)(Heavy Edgeの根から$k$個親にさかのぼる操作が出来ないと錯覚してしまった…)　いつか書き直す

使い方

• Tree(g, root = 0): rootを根としてデータ構造を構築する。(以下のクエリでgはrootを根とした根付き木とみなされる。)　$\mathrm{O}(N \log N)$
• depth(u): uの深さを返す。$\mathrm{O}(1)$
• par(u): uの親ノードを返す。(存在しない場合-1を返す。)　$\mathrm{O}(1)$
• kth_ancestor(u, k): $\mathrm{par}^k(u)$を返す。(存在しない場合-1を返す。)　$\mathrm{O}(\log N)$
• nxt(s, t): s-t間のパス上の点のうちsに隣接する点を返す。 $\mathrm{O}(\log N)$
• lca(u, v): uとvのlcaを返す。$\mathrm{O}(\log N)$
• path(s, t): s-t間のパス$p$をパスに所属する頂点の列で返す。$\mathrm{O}(\lvert p\rvert)$

Depends on

• グラフテンプレート (graph/graph-template.hpp)

Verified with

• verify/verify-unit-test/tree-path.test.cpp
• verify/verify-yosupo-graph/yosupo-jump-on-tree.test.cpp
• verify/verify-yosupo-graph/yosupo-lowest-common-ancestor-tree-util.test.cpp
• verify/verify-yuki/yuki-2588.test.cpp

Code

#pragma once
#include "../graph/graph-template.hpp"

template <typename G>
struct Tree {
 private:
  G& g;
  int root;
  vector<array<int, 24>> bl;
  vector<int> dp;
  void build() {
    bl.resize(g.size());
    dp.resize(g.size());
    for (auto& v : bl) fill(begin(v), end(v), -1);
    dfs(root, -1, 0);
  }

  void dfs(int c, int p, int _dp) {
    dp[c] = _dp;
    for (int i = p, x = 0; i != -1;) {
      bl[c][x] = i;
      i = bl[i][x], x++;
    }
    for (auto& d : g[c]) {
      if (d == p) continue;
      dfs(d, c, _dp + 1);
    }
  }

 public:
  Tree(G& _g, int _r = 0) : g(_g), root(_r) { build(); }

  int depth(int u) const { return dp[u]; }

  int par(int u) const { return u == root ? -1 : bl[u][0]; }

  int kth_ancestor(int u, int k) const {
    if (dp[u] < k) return -1;
    while (k) {
      int t = __builtin_ctz(k);
      u = bl[u][t], k ^= 1 << t;
    }
    return u;
  }

  int nxt(int s, int t) const {
    if (dp[s] >= dp[t]) return par(s);
    int u = kth_ancestor(t, dp[t] - dp[s] - 1);
    return bl[u][0] == s ? u : bl[s][0];
  }

  vector<int> path(int s, int t) const {
    vector<int> pre, suf;
    while (dp[s] > dp[t]) {
      pre.push_back(s);
      s = bl[s][0];
    }
    while (dp[s] < dp[t]) {
      suf.push_back(t);
      t = bl[t][0];
    }
    while (s != t) {
      pre.push_back(s);
      suf.push_back(t);
      s = bl[s][0];
      t = bl[t][0];
    }
    pre.push_back(s);
    reverse(begin(suf), end(suf));
    copy(begin(suf), end(suf), back_inserter(pre));
    return pre;
  }

  int lca(int u, int v) {
    if (dp[u] != dp[v]) {
      if (dp[u] > dp[v]) swap(u, v);
      v = kth_ancestor(v, dp[v] - dp[u]);
    }
    if (u == v) return u;
    for (int i = __lg(dp[u]); i >= 0; --i) {
      if (dp[u] < (1 << i)) continue;
      if (bl[u][i] != bl[v][i]) u = bl[u][i], v = bl[v][i];
    }
    return bl[u][0];
  }

  // u - v 間のパス上の頂点のうち u から距離 i の頂点
  // (dist(u, v) < i のとき -1)
  int jump(int u, int v, int i) {
    int lc = lca(u, v);
    int d1 = dp[u] - dp[lc];
    if (i <= d1) return kth_ancestor(u, i);
    int d = d1 + dp[v] - dp[lc];
    if (i <= d) return kth_ancestor(v, d - i);
    return -1;
  }
};

/**
 * @brief 木に対する一般的なクエリ
 * @docs docs/tree/tree-query.md
 */

#line 2 "graph/graph-template.hpp"

template <typename T>
struct edge {
  int src, to;
  T cost;

  edge(int _to, T _cost) : src(-1), to(_to), cost(_cost) {}
  edge(int _src, int _to, T _cost) : src(_src), to(_to), cost(_cost) {}

  edge &operator=(const int &x) {
    to = x;
    return *this;
  }

  operator int() const { return to; }
};
template <typename T>
using Edges = vector<edge<T>>;
template <typename T>
using WeightedGraph = vector<Edges<T>>;
using UnweightedGraph = vector<vector<int>>;

// Input of (Unweighted) Graph
UnweightedGraph graph(int N, int M = -1, bool is_directed = false,
                      bool is_1origin = true) {
  UnweightedGraph g(N);
  if (M == -1) M = N - 1;
  for (int _ = 0; _ < M; _++) {
    int x, y;
    cin >> x >> y;
    if (is_1origin) x--, y--;
    g[x].push_back(y);
    if (!is_directed) g[y].push_back(x);
  }
  return g;
}

// Input of Weighted Graph
template <typename T>
WeightedGraph<T> wgraph(int N, int M = -1, bool is_directed = false,
                        bool is_1origin = true) {
  WeightedGraph<T> g(N);
  if (M == -1) M = N - 1;
  for (int _ = 0; _ < M; _++) {
    int x, y;
    cin >> x >> y;
    T c;
    cin >> c;
    if (is_1origin) x--, y--;
    g[x].emplace_back(x, y, c);
    if (!is_directed) g[y].emplace_back(y, x, c);
  }
  return g;
}

// Input of Edges
template <typename T>
Edges<T> esgraph([[maybe_unused]] int N, int M, int is_weighted = true,
                 bool is_1origin = true) {
  Edges<T> es;
  for (int _ = 0; _ < M; _++) {
    int x, y;
    cin >> x >> y;
    T c;
    if (is_weighted)
      cin >> c;
    else
      c = 1;
    if (is_1origin) x--, y--;
    es.emplace_back(x, y, c);
  }
  return es;
}

// Input of Adjacency Matrix
template <typename T>
vector<vector<T>> adjgraph(int N, int M, T INF, int is_weighted = true,
                           bool is_directed = false, bool is_1origin = true) {
  vector<vector<T>> d(N, vector<T>(N, INF));
  for (int _ = 0; _ < M; _++) {
    int x, y;
    cin >> x >> y;
    T c;
    if (is_weighted)
      cin >> c;
    else
      c = 1;
    if (is_1origin) x--, y--;
    d[x][y] = c;
    if (!is_directed) d[y][x] = c;
  }
  return d;
}

/**
 * @brief グラフテンプレート
 * @docs docs/graph/graph-template.md
 */
#line 3 "tree/tree-query.hpp"

template <typename G>
struct Tree {
 private:
  G& g;
  int root;
  vector<array<int, 24>> bl;
  vector<int> dp;
  void build() {
    bl.resize(g.size());
    dp.resize(g.size());
    for (auto& v : bl) fill(begin(v), end(v), -1);
    dfs(root, -1, 0);
  }

  void dfs(int c, int p, int _dp) {
    dp[c] = _dp;
    for (int i = p, x = 0; i != -1;) {
      bl[c][x] = i;
      i = bl[i][x], x++;
    }
    for (auto& d : g[c]) {
      if (d == p) continue;
      dfs(d, c, _dp + 1);
    }
  }

 public:
  Tree(G& _g, int _r = 0) : g(_g), root(_r) { build(); }

  int depth(int u) const { return dp[u]; }

  int par(int u) const { return u == root ? -1 : bl[u][0]; }

  int kth_ancestor(int u, int k) const {
    if (dp[u] < k) return -1;
    while (k) {
      int t = __builtin_ctz(k);
      u = bl[u][t], k ^= 1 << t;
    }
    return u;
  }

  int nxt(int s, int t) const {
    if (dp[s] >= dp[t]) return par(s);
    int u = kth_ancestor(t, dp[t] - dp[s] - 1);
    return bl[u][0] == s ? u : bl[s][0];
  }

  vector<int> path(int s, int t) const {
    vector<int> pre, suf;
    while (dp[s] > dp[t]) {
      pre.push_back(s);
      s = bl[s][0];
    }
    while (dp[s] < dp[t]) {
      suf.push_back(t);
      t = bl[t][0];
    }
    while (s != t) {
      pre.push_back(s);
      suf.push_back(t);
      s = bl[s][0];
      t = bl[t][0];
    }
    pre.push_back(s);
    reverse(begin(suf), end(suf));
    copy(begin(suf), end(suf), back_inserter(pre));
    return pre;
  }

  int lca(int u, int v) {
    if (dp[u] != dp[v]) {
      if (dp[u] > dp[v]) swap(u, v);
      v = kth_ancestor(v, dp[v] - dp[u]);
    }
    if (u == v) return u;
    for (int i = __lg(dp[u]); i >= 0; --i) {
      if (dp[u] < (1 << i)) continue;
      if (bl[u][i] != bl[v][i]) u = bl[u][i], v = bl[v][i];
    }
    return bl[u][0];
  }

  // u - v 間のパス上の頂点のうち u から距離 i の頂点
  // (dist(u, v) < i のとき -1)
  int jump(int u, int v, int i) {
    int lc = lca(u, v);
    int d1 = dp[u] - dp[lc];
    if (i <= d1) return kth_ancestor(u, i);
    int d = d1 + dp[v] - dp[lc];
    if (i <= d) return kth_ancestor(v, d - i);
    return -1;
  }
};

/**
 * @brief 木に対する一般的なクエリ
 * @docs docs/tree/tree-query.md
 */

Back to top page