Skip to content

tree_diameter.hpp

#include "noya/tree_diameter.hpp"

View on GitHub

#ifndef NOYA_TREE_DIAMETER_HPP
#define NOYA_TREE_DIAMETER_HPP 1

#include "noya/lowest_common_ancestor.hpp"
#include "noya/shortest_path.hpp"

#include <algorithm>
#include <array>
#include <tuple>
#include <vector>

namespace noya {

/// @brief Unweighted tree diameter via double BFS. @return (diameter, u, v, eccentricity[]).
inline std::tuple<int, int, int, std::vector<int>>
tree_diam(std::vector<std::vector<int>> &g) {
  if (g.empty())
    return {-1, -1, -1, {}};
  auto d0 = bfs_unweighted(g, 0).first;
  int p = std::max_element(d0.begin(), d0.end()) - d0.begin();
  auto dp = bfs_unweighted(g, p).first;
  int q = std::max_element(dp.begin(), dp.end()) - dp.begin();
  auto dq = bfs_unweighted(g, q).first;
  int n = int(g.size());
  std::vector<int> ecc(n);
  for (int i = 0; i < n; i++)
    ecc[i] = std::max(dp[i], dq[i]);
  return {dp[q], p, q, ecc};
}

/// @brief Diameter monoid for segment tree. Merge two vertex sets and track the farthest pair.
struct diameter_monoid {
  using S = std::pair<int64_t, std::array<int, 2>>;
  static constexpr S identity = {-1, {-1, -1}};

  fastlca *lca;
  explicit diameter_monoid(fastlca &l) : lca(&l) {}

  S make(int v) const { return {0, {v, v}}; }

  S merge(S a, S b) const {
    if (a == identity) return b;
    if (b == identity) return a;
    S c = std::max(a, b);
    for (auto x : a.second)
      for (auto y : b.second) {
        int64_t d = lca->distance(x, y);
        if (d > c.first)
          c = {d, {x, y}};
      }
    return c;
  }
};

} // namespace noya

#endif // NOYA_TREE_DIAMETER_HPP
#include <algorithm>
#include <array>
#include <cassert>
#include <limits>
#include <queue>
#include <tuple>
#include <vector>

namespace noya {
template <class T, auto op, T e = T()> struct sparse_table {
  static int highest_bit(unsigned x) {
    return x == 0 ? -1 : 31 - __builtin_clz(x);
  }

  int n = 0;
  std::vector<std::vector<T>> ST;

  sparse_table(const std::vector<T> &values = {}) {
    if (!values.empty())
      build(values);
  }

  void build(const std::vector<T> &values) {
    n = int(values.size());
    int levels = highest_bit(n) + 1;
    ST.resize(levels);

    for (int k = 0; k < levels; k++)
      ST[k].resize(n - (1 << k) + 1);

    if (n > 0)
      ST[0] = values;

    for (int k = 1; k < levels; k++)
      for (int i = 0; i <= n - (1 << k); i++)
        ST[k][i] = op(ST[k - 1][i], ST[k - 1][i + (1 << (k - 1))]);
  }
  /// @brief Query op over [l, r).
  T prod(int l, int r) const {
    assert(0 <= l && l <= r && r <= n);
    if (l == r) {
      return e;
    }
    int L = highest_bit(r - l);
    return op(ST[L][l], ST[L][r - (1 << L)]);
  }
};
} // namespace noya

namespace noya {
/// @brief Sparse-table-based LCA with O(n log n) build and O(1) query.
struct fastlca {
  static std::pair<int, int> min_op(std::pair<int, int> a,
                                    std::pair<int, int> b) {
    return std::min(a, b);
  }
  int n;
  std::vector<int> dfn;
  std::vector<int> d;
  std::vector<int64_t> len;
  std::vector<int> siz;
  sparse_table<std::pair<int, int>, min_op> ST;
  bool weighted;

  std::vector<int> fa;

  template <class T>
  fastlca(const std::vector<T> &g = {}, const bool &_weighted = false,
          const int &root = 0) {
    weighted = _weighted;
    if (!g.empty())
      build(g, root);
  };

  template <class T>
  void build(const std::vector<std::vector<T>> &g, const int &root = 0) {
    n = int(g.size());
    std::vector<std::vector<int>> g2(n);

    for (int u = 0; u < n; u++) {
      for (auto &[v, w] : g[u]) {
        g2[u].push_back(v);
      }
    }
    build(g2, root);

    len.assign(n, 0);
    auto dfs = [&](auto &dfs, int u) -> void {
      for (auto &[v, w] : g[u]) {
        if (v == fa[u])
          continue;
        len[v] = len[u] + w;
        dfs(dfs, v);
      }
    };

    dfs(dfs, root);
  }

  void build(const std::vector<std::vector<int>> &g = {}, const int &root = 0) {
    n = int(g.size());
    d.assign(n, 0);
    dfn.assign(n, -1);
    siz.assign(n, 0);
    fa.assign(n, -1);

    int idx = 0;
    std::vector<std::pair<int, int>> a;
    a.reserve(n);

    auto dfs = [&](auto &dfs, int u, int p) -> void {
      fa[u] = p;
      siz[u] = 1;
      dfn[u] = idx++;
      if (p == -1)
        a.push_back({-1, -1});
      else
        a.push_back(std::make_pair(dfn[p], p));

      for (auto &v : g[u]) {
        if (v != p) {
          d[v] = d[u] + 1;
          dfs(dfs, v, u);
          siz[u] += siz[v];
        }
      }
    };

    dfs(dfs, root, -1);
    ST.build(a);
  }

  /// @brief Check if node a is in the subtree of node b.
  bool is_subtree(int a, int b) {
    if (dfn[b] <= dfn[a] && dfn[a] < dfn[b] + siz[b]) {
      return true;
    } else {
      return false;
    }
  }

  /// @brief Return the lowest common ancestor of nodes u and v.
  int lca(int u, int v) const {
    assert(0 <= u && u < n);
    assert(0 <= v && v < n);

    if (u == v) {
      return u;
    } else {
      int a = dfn[u];
      int b = dfn[v];
      if (a > b) {
        std::swap(a, b);
      }
      return ST.prod(a + 1, b + 1).second;
    }
  }

  /// @brief Return the LCA when the tree is re-rooted at c.
  int rooted_lca(int a, int b, int c) const {
    return lca(a, b) ^ lca(a, c) ^ lca(b, c);
  }

  /// @brief Return the (weighted or unweighted) distance between nodes a and b.
  int64_t distance(int a, int b) const {
    int c = lca(a, b);
    if (!weighted) {
      return d[a] + d[b] - d[c] * 2;
    } else {
      return len[a] + len[b] - len[c] * 2;
    }
  }

  /// @brief Compute the intersection of paths (a,b) and (c,d) as a pair of endpoints.
  std::pair<int, int> intersection(int a, int b, int c, int d) const {
    int ab = lca(a, b), ac = lca(a, c), ad = lca(a, d);
    int bc = lca(b, c), bd = lca(b, d), cd = lca(c, d);
    int x = ab ^ ac ^ bc;
    int y = ab ^ ad ^ bd;
    if (x != y) {
      return {x, y};
    }
    int z = ac ^ ad ^ cd;
    if (x != z) {
      x = -1;
    }
    return {x, x};
  }

  std::pair<int, int> intersection(std::pair<int, int> a,
                                   std::pair<int, int> b) const {
    return intersection(a.first, a.second, b.first, b.second);
  }

  /// @brief Return the list of vertices on the path from a to b.
  std::vector<int> path(int a, int b) {
    int c = lca(a, b);
    std::vector<int> ac;
    while (a != c) {
      ac.push_back(a);
      a = fa[a];
    }
    std::vector<int> bc;
    while (b != c) {
      bc.push_back(b);
      b = fa[b];
    }
    std::vector<int> res = std::move(ac);
    res.push_back(c);
    res.insert(res.end(), bc.rbegin(), bc.rend());
    return res;
  }
};
} // namespace noya

namespace noya {
/// @brief BFS on unweighted graph. @return (distances, predecessors).
inline std::pair<std::vector<int>, std::vector<int>>
bfs_unweighted(std::vector<std::vector<int>> &g, int start) {
  int N = int(g.size());
  const int INF = std::numeric_limits<int>::max();
  std::vector<int> dis(N, INF);
  std::vector<int> pre(N, -1);
  dis[start] = 0;
  pre[start] = start;

  std::vector<int> que{start};
  for (int i = 0; i < int(que.size()); i++) {
    int u = que[i];
    for (auto v : g[u])
      if (dis[v] == INF) {
        dis[v] = dis[u] + 1;
        pre[v] = u;
        que.push_back(v);
      }
  }
  return {dis, pre};
}

/// @brief Dijkstra's algorithm. @return (distances, predecessors).
template <class T>
std::pair<std::vector<int64_t>, std::vector<int>>
dijkstra(std::vector<std::vector<T>> &g, int start) {
  int N = int(g.size());
  const int64_t INF = std::numeric_limits<int64_t>::max();
  std::vector<int64_t> dis(N, INF);
  std::vector<int> pre(N, -1);
  std::priority_queue<std::pair<int64_t, int>,
                      std::vector<std::pair<int64_t, int>>, std::greater<>>
      que;
  que.emplace(dis[start] = 0, start);
  pre[start] = start;
  while (!que.empty()) {
    auto [d, u] = que.top();
    que.pop();
    if (d > dis[u])
      continue;
    for (auto &[v, w] : g[u])
      if (dis[v] > dis[u] + w) {
        dis[v] = dis[u] + w;
        pre[v] = u;
        que.emplace(dis[v], v);
      }
  }
  return {dis, pre};
}

inline std::vector<int> find_path(std::vector<int> &pre, int s, int t) {
  std::vector<int> path;
  int cur = t;
  while (cur != s) {
    assert(cur >= 0);
    path.push_back(cur);
    cur = pre[cur];
  }
  path.push_back(s);
  std::reverse(path.begin(), path.end());
  return path;
}

template <class T>
std::vector<std::vector<int64_t>> floyd(std::vector<std::vector<T>> &g) {
  int N = int(g.size());
  const int64_t INF = std::numeric_limits<T>::max() / 2;
  std::vector<std::vector<int64_t>> f(N, std::vector<int64_t>(N, INF));
  for (int i = 0; i < N; i++)
    for (int j = 0; j < N; j++)
      f[i][j] = g[i][j];
  for (int i = 0; i < N; i++)
    f[i][i] = 0;
  for (int k = 0; k < N; k++)
    for (int i = 0; i < N; i++)
      for (int j = 0; j < N; j++)
        f[i][j] = std::min(f[i][j], f[i][k] + f[k][j]);
  return f;
}

} // namespace noya

namespace noya {

/// @brief Unweighted tree diameter via double BFS. @return (diameter, u, v, eccentricity[]).
inline std::tuple<int, int, int, std::vector<int>>
tree_diam(std::vector<std::vector<int>> &g) {
  if (g.empty())
    return {-1, -1, -1, {}};
  auto d0 = bfs_unweighted(g, 0).first;
  int p = std::max_element(d0.begin(), d0.end()) - d0.begin();
  auto dp = bfs_unweighted(g, p).first;
  int q = std::max_element(dp.begin(), dp.end()) - dp.begin();
  auto dq = bfs_unweighted(g, q).first;
  int n = int(g.size());
  std::vector<int> ecc(n);
  for (int i = 0; i < n; i++)
    ecc[i] = std::max(dp[i], dq[i]);
  return {dp[q], p, q, ecc};
}

/// @brief Diameter monoid for segment tree. Merge two vertex sets and track the farthest pair.
struct diameter_monoid {
  using S = std::pair<int64_t, std::array<int, 2>>;
  static constexpr S identity = {-1, {-1, -1}};

  fastlca *lca;
  explicit diameter_monoid(fastlca &l) : lca(&l) {}

  S make(int v) const { return {0, {v, v}}; }

  S merge(S a, S b) const {
    if (a == identity) return b;
    if (b == identity) return a;
    S c = std::max(a, b);
    for (auto x : a.second)
      for (auto y : b.second) {
        int64_t d = lca->distance(x, y);
        if (d > c.first)
          c = {d, {x, y}};
      }
    return c;
  }
};

} // namespace noya