2021年4月4日日曜日

全方位木DP (Re-Rooting) の抽象的定式化と実装:頂点情報と辺情報を扱う

この記事では、全方位木DP(英語では re-rooting と呼ばれることが多いようです)と呼ばれるアルゴリズムを、できるだけ多くの問題を扱えるように抽象的に定式化しつつ、ライブラリとして使える実装を示すことを目的とします。全方位木DPを知らない方はまず keymoon さんによる【全方位木DP】明日使える便利な木構造のアルゴリズムをご覧下さい。 

全方位木DPによって解ける問題とは、どのように定式化できる問題でしょうか?

以下のように書ける問題は、全方位木DP によって解くことができます。無向木 \( T = (V, E) \) が与えられるとする。また、\( T \) を頂点 \( i \) を根とする根付き木として見たものを \( T_i \) とする。ここで、根付き木の部分木 \( S \) について次のように再帰的に定義される関数 \( F(S) \) を考えます。

$$  F(S) = \textrm{add_vertex} (G(S_1) \bullet \cdots \bullet G(S_m)), v) $$
$$  G(S_i) = \textrm{add_edge}(F(S_i), e_i) $$

但し
  • \( v \) は \( S \) の根
  • \( S_i \) は \( v \) の \( i \) 番目の子 \( v_i \) を根とする部分木
  • \( e_i \) は \( v \) と \( v_i \) を接続する辺
とします。




このとき \( F(T_i) \) をすべての頂点 \( i \) について求める問題が、全方位木DPによって解ける問題となります。

関数 \( F \) は次の三つの関数と一つの値によって定まることがわかります。辺情報を処理する関数 \( \textrm{add_edge} \)、子部分木の情報をマージするモノイド上の演算 \( \bullet \) とその単位元 \( \textrm{id} \)(モノイドについては上記の keymoon さんの記事を参照してください)、頂点情報を処理する関数 \( \textrm{add_vertex} \) です。従って、全方位木DPのアルゴリズムは、これら三つの関数及びモノイドの単位元を受け取るような関数として、ライブラリとして再利用できる形で実装することができます。実装は記事の最後に示します。

具体的な問題にあてはめることによって、様々な問題をこの定式化によって扱うことができることを見ていきましょう。

問題1: AOJ 木の直径 (問題リンク)
重み付き木における木の直径を求めてください。
この問題は全方位木DPを使わなくても解くことができますが、練習のため全方位木DPで解いてみましょう。

関数 \( F(S) \) を \( S \) の根から最も遠い頂点までの距離としましょう。\( \max(\{ F(T_i) | i \in V \}) \) が木の直径となります。

  • \( \textrm{add_edge}(x, e) = x + [e\textrm{'s length}] \)
  • \( x \bullet y = \max(x, y) \)
  • \( \textrm{id} = 0 \)
  • \( \textrm{add_vertex}(x, v) = x \)
とすると、関数 \( F(S) \) を定義することができます。

この問題では頂点情報は必要ないため \( \textrm{add_vertex} \) は自明な関数となっています。

ライブラリを使うと以下のように全方位木DPを実行することができます。
ReRooting<int64, int64>(
      graph,
      /* add_edge */ [](int64 x, const Edge e) -> int64 { return x + e.w; },
      /* monoid */ [](int64 x1, int64 x2) -> int64 { return max(x1, x2); },
      /* monoid id */ 0,
      /* add_vertex */ [](int64 x, int v) -> int64 { return x; });

問題2: AtCoder EDPC V Subtree問題リンク
木が与えられます。各頂点を白または黒に塗りますが、黒い頂点が連結になるようにします。このとき、各頂点 \( v \) が黒になるような塗り方は何通りあるかを \( \mod M \) で求めてください。
関数 \( F(S) \) を「部分木 \( S \) の根を黒く塗り、かつ黒い頂点が連結になるような色の塗り方の数」とします。各頂点 \( v \) に対して \( F(T_v) \) が求める値となります。
  • \( \textrm{add_edge}(x, e) = x + 1 \)
  • \( x \bullet y = x \cdot y \)
  • \( \textrm{id} = 1 \)
  • \( \textrm{add_vertex}(x, v) = x \)
とすると、関数 \( F(S) \) を定義することができます(詳細な解説は kyopro_friends さんによる解説を参照して下さい)。

ライブラリを使うと以下のように全方位木DPを実行することができます。
ReRooting<ModInt, ModInt>(
     graph,
      /* add_edge */ [](ModInt x, Edge e) -> { return x + 1; },
      /* monoid */ [](ModInt x, ModInt y) -> { return x * y; },
      /* monoid id */ 1,
      /* add_vertext */ [](ModInt x, int v) -> { return x; });
TODO: 頂点情報を使う問題を発見したら例として足す

ライブラリの実装例
struct Edge {
  const int from, to;
  Edge(int from, int to) : from(from), to(to) {}
};

template<class T, class U> vector<T> ReRooting(
    const vector<vector<Edge>>& graph, function<U(T, Edge)> edge_fn,
    function<U(U, U)> merge_fn, U merge_id, function<T(U, int)> vertex_fn) {
  int n = graph.size();

  // BFS
  vector<int> vs({0}), ps(n, -2); ps[0] = -1; queue<int> q; q.push(0);
  while (!q.empty()) {
    int v = q.front(); q.pop();
    for (const auto& e : graph[v]) {
      if (e.to == ps[e.to] || ps[e.to] >= -1) continue;
      vs.push_back(e.to); ps[e.to] = v; q.push(e.to);
    }
  }

  // DP phase 1.
  vector<T> dp_1(n);
  for (int i = n - 1; i >= 0; i--) {
    int v = vs[i]; int p = ps[v];
    auto a = merge_id;
    for (const auto& e : graph[v]) {
      if (e.to == ps[v]) continue; a = merge_fn(a, edge_fn(dp_1[e.to], e));
    }
    dp_1[v] = vertex_fn(a, v);
  }

  // DP phase 2.
  vector<T> dp_2(n), dp_3(n); dp_2[0] = merge_id;
  for (int i = 0; i < n; i++) {
    int v = vs[i]; int p = ps[v]; int m = graph[v].size();
    vector<T> ls(m + 1), rs(m + 1);
    ls[0] = merge_id;
    for (int j = 0; j < m; j++) {
      const auto& e = graph[v][j];
      if (e.to == p) ls[j + 1] = ls[j];
      else ls[j + 1] = merge_fn(ls[j], edge_fn(dp_1[e.to], e));
    }
    rs[m] = merge_id;
    for (int j = m - 1; j >= 0; j--) {
      const auto& e = graph[v][j];
      if (e.to == p) rs[j] = rs[j + 1];
      else rs[j] = merge_fn(rs[j + 1], edge_fn(dp_1[e.to], e));
    }
    dp_3[v] = vertex_fn(merge_fn(ls[m], dp_2[v]), v);
    for (int j = 0; j < m; j++) {
      const auto& e = graph[v][j]; if (e.to == p) continue;
      for (const auto& re : graph[e.to]) {
        if (re.to != v) continue;
        T merged = merge_fn(merge_fn(ls[j], rs[j + 1]), dp_2[v]);
        dp_2[e.to] = edge_fn(vertex_fn(merged, v), re);
      }
    }
  }
  return dp_3;
}

HTTF 2023 予選 (AHC 016) 3位解法の解説

2022年11月に開催された HACK TO THE FUTURE 2023 予選 (AtCoder Heuristic Contest 016) の3位解法の解説をします。 \( \epsilon = 0 \) の場合 この場合は厳密に解けるので、特別扱いします。頂点数毎の非...