Rerooting DP
(tree/RerootingDP.hpp)

View this file on GitHub
Last update: 2023-05-17 22:41:14+09:00
Include: #include "tree/RerootingDP.hpp"

RerootingDP

木に対して全方位木DPを行います。次の問題を解きます。

頂点 $0,1,\dots,n-1$ からなる $n$ 頂点の木 $T$ が与えられます。有向グラフ $G=(\mathbb{V},\mathbb{E})$ は以下を満たします。

$\mathbb{V}=\lbrace 0,1,\dots, n-1 \rbrace$ である。

$\mathbb{E}\subset \mathbb{V}\times \mathbb{V}$ である。辺 $(a,b)$ は頂点 $a$ から頂点 $b$ への有向辺。 $T$ において頂点 $a,b$ の間に辺が張られているとき, またそのときに限り $(a,b),(b,a)\in \mathbb{E}$ である。

また、各有向辺には整数値の番号が割り当てられていて、 $\mathrm{idx}:\mathbb{E}\to \mathbb{Z}$ によって定められます。ここで $\mathbb{E}_\mathrm{idx}=\lbrace \mathrm{idx}(e)\ \mid\ e\in \mathbb{E} \rbrace$ としておきます。

可換モノイド $(E,\oplus : E\times E\to E,e\in E)$ と集合 $V$ および $\mathrm{put _ edge}:V\times \mathbb{E}_\mathrm{idx}\to E,\ \mathrm{put _ vertex}:E\times \mathbb{V}\to V$ が与えられます。次の擬似コードに従って計算される値 dfs(0,0), dfs(1,1),...,dfs(n-1,n-1) を求めてください。ただし、頂点 $r$ を根と見たときの頂点 $v$ の子の頂点の集合を $\mathrm{child}(r,v)$ とします。
V dfs(vertex r, vertex v){
   E prod = e();
   for ( (v, c) in E, c in child(r,v) ){
       prod = merge(prod, put_edge(dfs(r,c),idx(v,c)));
   }
   return put_vertex(prod, v);
}

この問題を時間計算量 $\mathrm{O}(n)$ で解くことができます。

また、このライブラリはオラクルとして merge, e, put_edge, put_vertex を使用しますが、これらが定数時間で動くものと仮定したときの計算量を記述します。オラクル内部の計算量が $\mathrm{O}(f(n))$ である場合はすべての計算量が $\mathrm{O}(f(n))$ 倍となります。

コンストラクタ

RerootingDP<E, V, merge, e, put_edge, put_vertex> g(int n);

可換モノイドの型 E
DPの値の型 V
$\oplus : E\times E\to E$ を計算する関数 E merge(E a, E b)
$e$ を返す関数 E e()
辺番号 $i$ の辺を付加する関数 E put_edge(V x, int i)
頂点番号 $v$ の頂点を付加する関数 V put_vertex(E x, int v)

を定義する必要があります。

頂点数 n の木を作る準備をします。n 頂点 $0$ 辺のグラフを作ります。
デフォルトでは n = 0 です。

制約

$0\le n\le 10^8$

計算量

$\mathrm{O}(n)$

add_edge

void g.add_edge(int u, int v, int idx1, int idx2)

頂点 u から頂点 v への辺番号が idx1 である辺を追加します。
頂点 v から頂点 u への辺番号が idx2 である辺を追加します。

制約

$0\le u,v\lt n$
ちょうど $n-1$ 回 add_edge を呼んでできるグラフは連結（どの $2$ 頂点間にもパスが存在する）であり、その後 add_edge は呼ばれない。

計算量

ならし $\mathrm{O}(1)$

build

vector<V> g.build(int r)

長さ $n$ の配列 subdp を返します。subdp[v] は上述した問題における dfs(r,r) を計算する過程で dfs が返す値 dfs(r,v) です。
頂点 r を根とする木DPのすべての結果を取得することができます。

制約

build を呼ぶ前にちょうど $n-1$ 回 add_edge を呼んでいる。
build は $1$ 回までしか呼ばれない。

計算量

$\mathrm{O}(n)$

reroot

vector<V> g.reroot()

長さ $n$ の配列 answers を返します。answers[v] は上述した問題における dfs(v,v) です。

制約

reroot を呼ぶ前にちょうど $1$ 回 bulid を呼んでいる。
reroot は $1$ 回までしか呼ばれない。

計算量

$\mathrm{O}(n)$

get

V get g.get(int r, int v)

上述した問題における dfs(r,r) を計算する過程で dfs が返す値 dfs(r,v) を返します。
頂点 r を根とする木DPをしたときの頂点 v の部分木に対するDPの値を得ることができます。

制約

get を呼ぶ前にちょうど $1$ 回 reroot を呼んでいる。

計算量

$\mathrm{O}(\log n)$

使用例

AC code of https://atcoder.jp/contests/dp/tasks/dp_v

submission https://atcoder.jp/contests/dp/submissions/36003936

#include<bits/stdc++.h>
#include<atcoder/modint>
using namespace std;
using mint = atcoder::modint;

#include"RerootingDP"

mint merge(mint a, mint b){
    return a * b;
}
mint e(){
    return mint(1);
}
mint put_edge(mint v, int i){
    return v + 1;
}
mint put_vertex(mint e, int v){
    return e;
}

int main(){
    int n, m; cin >> n >> m;
    mint::set_mod(m);
    RerootingDP<mint,mint,merge,e,put_edge,put_vertex> g(n);
    for (int i = 0; i < n-1; i++){
        int u, v; cin >> u >> v;
        g.add_edge(u-1,v-1,i,i);
    }
    g.build();
    for (auto ans : g.reroot()) cout << ans.val() << endl;
}

設計の意図

全方位木DPはすべての部分木 $t$ に対する答え $dp(t)$ を高速に求めるアルゴリズムで、 $t$ に含まれるさらに小さい部分木 $s$ に対する答え $dp(s)$ から $dp(t)$ を計算していきます。 $r$ を全体の根としたときの木DPの値を求める際、$v$ の部分木を有向辺 $(par(v),v)$ に対応させると理解しやすい、という話があります。（ $par(v)$ は $r$ を根としたときの $v$ の親です。 $v=r$ のときは $par(v)$ は仮想的な頂点であり、有向辺 $(par(v),v)$ は仮想的な辺です。）これに基づいて、有向辺 $(u,v)$ に対応する答えを $f(u,v)$ と書くことにしましょう。 $u=-1$ のときは仮想的な有向辺を考えており、対応する部分木は $v$ を根とする根付き木です。さて、このライブラリの設計では、

\[f(u,v)=\bigoplus_{c\in \mathrm{child}(u,v)}f(v,c)\]

のような遷移式では考慮されない「辺の重み」を取り入れるため、補助的な関数 $g(u,v)$ を定義しています。 $g(u,v)$ は $(u,v)$ に対応する部分木に辺 $(u,v)$ を付加するが頂点 $u$ はまだ付加しないみたいな（雰囲気で感じ取ってください）部分木もどきに対する答えっぽいもの（雰囲気で感じ取ってください）です。要点は次のとおりです。

部分木に辺を付加すると、部分木もどきになる。( $\mathrm{put \_ edge}$ )
部分木もどきをマージして頂点を付加すると、部分木になる。( $\mathrm{put \_ vertex}$ )
$\displaystyle f(u,v)=\mathrm{put _ vertex}\left(\bigoplus_{c\in \mathrm{child}(u,v)} \mathrm{put _ edge}(f(v,c),\mathrm{idx}(v,c)) ,\ v\right)$

ライブラリ内で辺に重み E e 、頂点に重み V v を持つような設計もあり得ると思ったのですが、できるだけライブラリの外側で問題依存パートを扱いたいのと、 put_vertex には単位元のようなものを要求しないため初期値の設定に困ったので、頂点・辺番号を引数に渡す関数を要求するライブラリの設計にしました。

Depends on

テンプレート (template/template.hpp)

Verified with

Code

#pragma once

#include "../template/template.hpp"

namespace lib {

template <class E, class V, E (*merge)(E, E), E (*e)(), E (*put_edge)(V, int),
          V (*put_vertex)(E, int)>
struct RerootingDP {
    struct edge {
        int to, idx, xdi;
    };
    RerootingDP(int _n = 0) : n(_n) {
        es.resize(n);
    }
    void add_edge(int u, int v, int idx1, int idx2) {
        es[u].push_back({v, idx1, idx2});
        es[v].push_back({u, idx2, idx1});
    }
    vector<V> build(int v = 0) {
        root = v;
        outs.resize(n);
        subdp.resize(n);
        in.resize(n), up.resize(n);
        int tnow = 0;
        dfs(root, -1, tnow);
        return subdp;
    }
    vector<V> reroot() {
        reverse_edge.resize(n);
        reverse_edge[root] = e();
        reverse_dp.resize(n);
        answers.resize(n);
        bfs(root);
        return answers;
    }
    V get(int r, int v) {
        if (r == v) return answers[r];
        if (!(in[v] < in[r] && up[r] <= up[v])) return subdp[v];
        int le = 0, ri = outs[v].size();
        while (ri - le > 1) {
            int md = (le + ri) / 2;
            if (in[es[v][md].to] <= in[r])
                le = md;
            else
                ri = md;
        }
        return reverse_dp[es[v][le].to];
    }
    const vector<edge> &operator[](int idx) const {
        return es[idx];
    }

  private:
    int n, root;
    vector<vector<edge>> es;
    vector<vector<E>> outs;
    vector<E> reverse_edge;
    vector<V> subdp, reverse_dp, answers;
    vector<int> in, up;
    void dfs(int v, int p, int &t) {
        E val = e();
        in[v] = t++;
        for (auto &u : es[v]) {
            if (u.to == p && u.to != es[v].back().to) swap(u, es[v].back());
            if (u.to == p) continue;
            dfs(u.to, v, t);
            E nval = put_edge(subdp[u.to], u.idx);
            outs[v].emplace_back(nval);
            val = merge(val, nval);
        }
        subdp[v] = put_vertex(val, v);
        up[v] = t;
    }
    void bfs(int v) {
        int siz = outs[v].size();
        vector<E> lui(siz + 1), rui(siz + 1);
        lui[0] = e(), rui[siz] = e();
        for (int i = 0; i < siz; i++) lui[i + 1] = merge(lui[i], outs[v][i]);
        for (int i = siz - 1; i >= 0; i--)
            rui[i] = merge(outs[v][i], rui[i + 1]);
        for (int i = 0; i < siz; i++) {
            reverse_dp[es[v][i].to] = put_vertex(
                merge(merge(lui[i], rui[i + 1]), reverse_edge[v]), v);
            reverse_edge[es[v][i].to] =
                put_edge(reverse_dp[es[v][i].to], es[v][i].xdi);
            bfs(es[v][i].to);
        }
        answers[v] = put_vertex(merge(lui[siz], reverse_edge[v]), v);
    }
};

}  // namespace lib

#line 2 "tree/RerootingDP.hpp"

#line 2 "template/template.hpp"

#include <bits/stdc++.h>

#define rep(i, s, n) for (int i = (int)(s); i < (int)(n); i++)
#define rrep(i, s, n) for (int i = (int)(n)-1; i >= (int)(s); i--)
#define all(v) v.begin(), v.end()

using ll = long long;
using ld = long double;
using ull = unsigned long long;

template <typename T> bool chmin(T &a, const T &b) {
    if (a <= b) return false;
    a = b;
    return true;
}
template <typename T> bool chmax(T &a, const T &b) {
    if (a >= b) return false;
    a = b;
    return true;
}

namespace lib {

using namespace std;

}  // namespace lib

// using namespace lib;
#line 4 "tree/RerootingDP.hpp"

namespace lib {

template <class E, class V, E (*merge)(E, E), E (*e)(), E (*put_edge)(V, int),
          V (*put_vertex)(E, int)>
struct RerootingDP {
    struct edge {
        int to, idx, xdi;
    };
    RerootingDP(int _n = 0) : n(_n) {
        es.resize(n);
    }
    void add_edge(int u, int v, int idx1, int idx2) {
        es[u].push_back({v, idx1, idx2});
        es[v].push_back({u, idx2, idx1});
    }
    vector<V> build(int v = 0) {
        root = v;
        outs.resize(n);
        subdp.resize(n);
        in.resize(n), up.resize(n);
        int tnow = 0;
        dfs(root, -1, tnow);
        return subdp;
    }
    vector<V> reroot() {
        reverse_edge.resize(n);
        reverse_edge[root] = e();
        reverse_dp.resize(n);
        answers.resize(n);
        bfs(root);
        return answers;
    }
    V get(int r, int v) {
        if (r == v) return answers[r];
        if (!(in[v] < in[r] && up[r] <= up[v])) return subdp[v];
        int le = 0, ri = outs[v].size();
        while (ri - le > 1) {
            int md = (le + ri) / 2;
            if (in[es[v][md].to] <= in[r])
                le = md;
            else
                ri = md;
        }
        return reverse_dp[es[v][le].to];
    }
    const vector<edge> &operator[](int idx) const {
        return es[idx];
    }

  private:
    int n, root;
    vector<vector<edge>> es;
    vector<vector<E>> outs;
    vector<E> reverse_edge;
    vector<V> subdp, reverse_dp, answers;
    vector<int> in, up;
    void dfs(int v, int p, int &t) {
        E val = e();
        in[v] = t++;
        for (auto &u : es[v]) {
            if (u.to == p && u.to != es[v].back().to) swap(u, es[v].back());
            if (u.to == p) continue;
            dfs(u.to, v, t);
            E nval = put_edge(subdp[u.to], u.idx);
            outs[v].emplace_back(nval);
            val = merge(val, nval);
        }
        subdp[v] = put_vertex(val, v);
        up[v] = t;
    }
    void bfs(int v) {
        int siz = outs[v].size();
        vector<E> lui(siz + 1), rui(siz + 1);
        lui[0] = e(), rui[siz] = e();
        for (int i = 0; i < siz; i++) lui[i + 1] = merge(lui[i], outs[v][i]);
        for (int i = siz - 1; i >= 0; i--)
            rui[i] = merge(outs[v][i], rui[i + 1]);
        for (int i = 0; i < siz; i++) {
            reverse_dp[es[v][i].to] = put_vertex(
                merge(merge(lui[i], rui[i + 1]), reverse_edge[v]), v);
            reverse_edge[es[v][i].to] =
                put_edge(reverse_dp[es[v][i].to], es[v][i].xdi);
            bfs(es[v][i].to);
        }
        answers[v] = put_vertex(merge(lui[siz], reverse_edge[v]), v);
    }
};

}  // namespace lib

Rerooting DP (tree/RerootingDP.hpp)