全方位木DPライブラリを作った

概要

森に対して全方位木DPを行うライブラリです。
全方位森DPの方がいいんかな。
森でないときにassertが落ちるようにするためにUnionFindを使っています。
パフォーマンスが気になるなら使わないように実装しなおすと良いかもしれません。

計算量

  • construct: O(n)
  • add: O(1)
  • dp: O(n\alpha(n)) (実質  O(n))

コード

using ll = long long;

class UnionFind {
  vector<ll> par, h;
public:
  UnionFind(ll size) : par(size, 0), h(size, 0) {
    for (int i = 0; i < size; ++i) {
      par[i] = i;
    }
  }
  void unite(ll u, ll v) {
    u = root(u), v = root(v);
    if (u == v) return;
    if (h[u] < h[v]) {
      par[u] = v;
    }
    else {
      par[v] = u;
    }
    if (h[u] == h[v]) ++h[u];
  }
  bool isUnited(ll u, ll v) {
    return root(u) == root(v);
  }
  ll root(ll v) {
    if (par[v] == v) return v;
    return par[v] = root(par[v]);
  }
};

template<class Calc>
class EdgeDpOnForest {
public:
  using T = typename Calc::type;
  struct Edge {
    ll to, rev;
    T value;
  };
private:
  ll n;
  void dfs1(ll v, ll pv) {
    ll idx = -1;
    vector<T> values;
    for (int i = 0; i < G[v].size(); ++i) {
      const Edge& e = G[v][i];
      if (e.to == pv) {
        idx = i;
        continue;
      }
      dfs1(e.to, v);
      values.push_back(e.value);
    }
    // 根以外では親から入る辺の値を更新
    if (idx >= 0) {
      const Edge& e = G[v][idx];
      G[e.to][e.rev].value = Calc::merge(values);
    }
  }
  void dfs2(ll v, ll pv) {
    vector<T> values;
    for (auto&& e : G[v]) values.push_back(e.value);
    values = Calc::evaluate(values);
    for (int i = 0; i < G[v].size(); ++i) {
      const Edge& e = G[v][i];
      if (e.to == pv) continue;
      G[e.to][e.rev].value = values[i];
      dfs2(e.to, v);
    }
  }
public:
  UnionFind uf;
  vector<vector<Edge>> G;
  EdgeDpOnForest(ll n) : n(n), uf(n), G(n) {}
  // 辺を追加する
  void add(ll u, ll v) {
    assert(!uf.isUnited(u, v));
    G[u].push_back({v, ll(G[v].size())});
    G[v].push_back({u, ll(G[u].size())-1});
    uf.unite(u, v);
  }
  // 各有向辺の値 G[v][i].value を計算する
  void dp() {
    vector<bool> used(n, false);
    for (int i = 0; i < n; ++i) {
      if (used[uf.root(i)]) continue;
      dfs1(i, -1);
      used[uf.root(i)] = true;
    }
    used.assign(n, false);
    for (int i = 0; i < n; ++i) {
      if (used[uf.root(i)]) continue;
      dfs2(i, -1);
      used[uf.root(i)] = true;
    }
  }
  size_t size() const {
    return G.size();
  }
};

// これらを実装してください
/*
struct Merger {
  using type = ll;
  // その頂点から出ていく値から入ってくる値を計算する
  static vector<type> evaluate(const vector<type>& value) {
  }
  // (辺数-1)個の出ていく値から1個の入ってくる値を計算する
  static type merge(const vector<type>& value) {
  }
};
*/

使用例 (問題: V - Subtree)

#include <bits/stdc++.h>

// ここに↑を貼る

ll M;

struct Merger {
  using type = ll;
  // その頂点から出ていく値から入ってくる値を計算する
  static vector<type> evaluate(const vector<type>& value) {
    const ll n = value.size();
    vector<type> L(n+1, 1), R(n+1, 1);
    for (int i = 0; i < n; ++i) {
      L[i+1] = (L[i] * value[i]) % M;
    }
    for (int i = n-1; i >= 0; --i) {
      R[i] = (R[i+1] * value[i]) % M;
    }
    vector<type> res(n);
    for (int i = 0; i < n; ++i) {
      res[i] = (L[i] * R[i+1] + 1) % M;
    }
    return res;
  }
  // (辺数-1)個の出ていく値から1個の入ってくる値を計算する
  static type merge(const vector<type>& value) {
    ll res = 1;
    for (auto&& x : value) {
      res = (res * x) % M;
    }
    res = (res + 1) % M;
    return res;
  }
};

int main() {
  ll n; cin >> n >> M;
  EdgeDpOnForest<Merger> dp(n);
  for (int i = 0; i < n-1; ++i) {
    ll a, b; cin >> a >> b; --a, --b;
    dp.add(a, b);
  }
  dp.dp();
  for (int i = 0; i < n; ++i) {
    ll ans = 1;
    for (auto&& e : dp.G[i]) {
      ans = (ans * e.value) % M;
    }
    cout << ans << endl;
  }
}