全方位木DP
(Algorithm/treedp.hpp)

• View this file on GitHub
• Last update: 2024-10-22 05:36:02+09:00
• Include: #include "Algorithm/treedp.hpp"

概要

抽象化した全方位木DP。逆元を仮定していない。

注意喚起

これを貼るだけで解けないのなら、ライブラリを改造して対応しようとせずに、一から書いた方が遥かに速い場合が多い。

全方位と相性が悪いであろう木DPのタイプ

• mergeに時間がかかるdp 例えば、2乗の木DP。全方位にせずに、全ての頂点について毎回木DPした方が良さそう。

このテンプレと相性が悪いタイプ

• mergeに補助的なデータ構造が必要な木dp
• 逆元を用いた場合に高速なタイプ 例: 子の頂点のうち、k番目に大きい値を根のスコアとする

手を加えなくても解ける問題のタイプ

• 部分木の結果のmergeが高速であるdp affine変換等

コンストラクタ

TDP<S, op, e, addroot, F, mp> tdp(n)

• S 部分木の結果を格納する型

• F 部分木の根に辺を追加することによる、部分木の結果への影響を計算するのに必要な情報を格納する型

• S op(l S, r S) 部分木の結果をmergeする関数。この際、l, r共に辺の情報は反映済みである。

• S e() 結果の単位元。単位元が思い浮かばない場合、返り値を特殊なものにしてopで弾く。

• S addroot(S s, int now) vの子の結果を(辺の情報を含めて)mergeしたものSと、自分の頂点番号vが与えられるのでそれにvの情報を付加し、vの結果を返す

• S mp(F f, S s) ある部分木の結果sに、辺の情報fを付加する。

• 計算量 $O(n)$

関数

• void add_edge(int s, int t, F f, F h)…s, t間に辺を張る。 s->tの重みがf, t->sの重みがh。
  • 制約 $0 \le s, t < n$
  • 計算量 $O(1)$
• vec<S> exe()…全方位木DPをする。
  • 計算量 addroot, mp, opをそれぞれ$O(頂点数)$回行う。

Verified with

• verify/treedp.test.cpp

Code

template <class S,
          S (*op)(S, S),
          S (*e)(),
          S (*addroot)(S, int),
          class F,
          S (*mp)(F, S)>
struct TDP {
    using pif = pair<int, F>;
    using vs = vec<S>;
    using vvs = vec<vs>;

    int n;
    vec<vec<pif>> g;
    vvs dp;
    // dp[v][i] :=  (v → g[v][i])の辺について、
    // g[v][i]を根とする部分木の結果
    vs ans;

    TDP(int n) : n(n) {
        g.resize(n);
        dp = vvs(n);
        ans = vs(n, e());
    }

  private:
    S dfs(int v, int p) {
        S res = e();
        int d = g[v].size();
        dp[v].resize(d);
        rep(i, 0, d) {
            int to = g[v][i].first;
            if (to == p) continue;
            dp[v][i] = dfs(to, v);
            res = op(res, mp(g[v][i].second, dp[v][i]));
            // 部分木の結果を集約。
            // 本実装では辺を加味 ->
            // 部分木集約の順を徹底している(辻褄が合うならいつでも良い)
        }
        // 辺・頂点をaddした影響を反映したものを返す。
        return addroot(res, v);
    }

    void bfs(int v, S par, int p) {
        int d = g[v].size();
        rep(i, 0, d) if (g[v][i].first == p) dp[v][i] = par;
        // 親の結果を渡しておく。

        vs L(d + 1, e());
        vs R(d + 1, e());

        rep(i, 0, d) L[i + 1] = op(L[i], mp(g[v][i].second, dp[v][i]));
        rrep(i, 0, d) R[i] = op(mp(g[v][i].second, dp[v][i]), R[i + 1]);
        // 本実装では辺を加味 ->
        // 部分木集約の順を徹底している(辻褄が合うならいつでも良い)

        ans[v] = addroot(L[d], v);
        // 辺・頂点をaddした影響を反映したものを返す。ansに格納する時だけ何か弄りたいならここを弄る。

        rep(i, 0, d) {
            int to = g[v][i].first;
            if (to == p) continue;
            S nx = op(L[i], R[i + 1]);
            // 本実装では辺を加味 ->
            // 部分木集約の順を徹底している(辻褄が合うならいつでも良い)

            bfs(to, addroot(nx, v), v);
            // to -> vの向きに辺・頂点をaddした影響を反映したものを返す。
        }
    }

  public:
    // s -> t に重みfの辺、 t -> sに重みhの辺
    void add_edge(int s, int t, F f, F h) {
        g[s].emplace_back(t, f);
        g[t].emplace_back(s, h);
    }

    vec<S> exe() {
        dfs(0, -1);
        bfs(0, e(), -1);
        return ans;
    }
};

/*
@brief 全方位木DP
@docs doc/treedp.md
*/

#line 1 "Algorithm/treedp.hpp"
template <class S,
          S (*op)(S, S),
          S (*e)(),
          S (*addroot)(S, int),
          class F,
          S (*mp)(F, S)>
struct TDP {
    using pif = pair<int, F>;
    using vs = vec<S>;
    using vvs = vec<vs>;

    int n;
    vec<vec<pif>> g;
    vvs dp;
    // dp[v][i] :=  (v → g[v][i])の辺について、
    // g[v][i]を根とする部分木の結果
    vs ans;

    TDP(int n) : n(n) {
        g.resize(n);
        dp = vvs(n);
        ans = vs(n, e());
    }

  private:
    S dfs(int v, int p) {
        S res = e();
        int d = g[v].size();
        dp[v].resize(d);
        rep(i, 0, d) {
            int to = g[v][i].first;
            if (to == p) continue;
            dp[v][i] = dfs(to, v);
            res = op(res, mp(g[v][i].second, dp[v][i]));
            // 部分木の結果を集約。
            // 本実装では辺を加味 ->
            // 部分木集約の順を徹底している(辻褄が合うならいつでも良い)
        }
        // 辺・頂点をaddした影響を反映したものを返す。
        return addroot(res, v);
    }

    void bfs(int v, S par, int p) {
        int d = g[v].size();
        rep(i, 0, d) if (g[v][i].first == p) dp[v][i] = par;
        // 親の結果を渡しておく。

        vs L(d + 1, e());
        vs R(d + 1, e());

        rep(i, 0, d) L[i + 1] = op(L[i], mp(g[v][i].second, dp[v][i]));
        rrep(i, 0, d) R[i] = op(mp(g[v][i].second, dp[v][i]), R[i + 1]);
        // 本実装では辺を加味 ->
        // 部分木集約の順を徹底している(辻褄が合うならいつでも良い)

        ans[v] = addroot(L[d], v);
        // 辺・頂点をaddした影響を反映したものを返す。ansに格納する時だけ何か弄りたいならここを弄る。

        rep(i, 0, d) {
            int to = g[v][i].first;
            if (to == p) continue;
            S nx = op(L[i], R[i + 1]);
            // 本実装では辺を加味 ->
            // 部分木集約の順を徹底している(辻褄が合うならいつでも良い)

            bfs(to, addroot(nx, v), v);
            // to -> vの向きに辺・頂点をaddした影響を反映したものを返す。
        }
    }

  public:
    // s -> t に重みfの辺、 t -> sに重みhの辺
    void add_edge(int s, int t, F f, F h) {
        g[s].emplace_back(t, f);
        g[t].emplace_back(s, h);
    }

    vec<S> exe() {
        dfs(0, -1);
        bfs(0, e(), -1);
        return ans;
    }
};

/*
@brief 全方位木DP
@docs doc/treedp.md
*/

Back to top page