Rerooting DP - 全方位木DP
(Library/Tree/RerootingDP.hpp)

View this file on GitHub
Last update: 2026-03-04 11:13:30+09:00
Include: #include "Library/Tree/RerootingDP.hpp"

Attention: このドキュメントは未完成です。

Rerooting DP - 全方位木DP

木のすべての頂点を根としたときの DP の値を効率的に計算するデータ構造です。

各頂点を根とした場合の DP を愚直に計算すると $\textrm{O}(N^2)$ かかりますが、全方位木 DP を用いることで $\textrm{O}(N)$ で計算できます。

Function

Constructor (辺属性のみ)

RerootingDP(Graph<CostType> &tree, Fsub merge, Gsub add, const Monoid monoid_identity, Vertex r = 0)

木 tree に対して全方位木 DP を行います（辺属性のみを考慮）。
merge: (Monoid, Monoid) -> Monoid - モノイド同士の二項演算
add: (Monoid, CostType) -> Monoid - モノイドと辺コストの二項演算
monoid_identity: モノイドの単位元
r: 初期の根（デフォルトは頂点 0）

型定義

Fsub = function<Monoid(Monoid, Monoid)>
Gsub = function<Monoid(Monoid, CostType)>

制約

$1 \le N \le 10^5$
merge と add はモノイドの演算を満たす

計算量

$\textrm{O}(N)$

Constructor (頂点属性を含む)

RerootingDP(Graph<CostType> &tree, F merge, G add, H finalize, const Monoid monoid_identity, Vertex r = 0)

木 tree に対して全方位木 DP を行います（頂点属性を含む）。
merge: (Monoid, Monoid, Vertex) -> Monoid - モノイド同士の二項演算（頂点情報を利用）
add: (Monoid, CostType, Vertex) -> Monoid - モノイドと辺コストの二項演算（頂点情報を利用）
finalize: (Monoid, Vertex) -> Monoid - 頂点が根のときの最終処理
monoid_identity: モノイドの単位元
r: 初期の根（デフォルトは頂点 0）

型定義

F = function<Monoid(Monoid, Monoid, Vertex)>
G = function<Monoid(Monoid, CostType, Vertex)>
H = function<Monoid(Monoid, Vertex)>

制約

$1 \le N \le 10^5$
merge, add, finalize はモノイドの演算を満たす

計算量

$\textrm{O}(N)$

GetAllAnswer

vector<Monoid> &GetAllAnswer()

全頂点に関する DP の結果の配列を取得します。
配列の $i$ 番目の要素は、頂点 $i$ を根としたときの DP の値です。

計算量

$\textrm{O}(1)$

operator[]

Monoid operator[](Vertex v)
const Monoid operator[](Vertex v) const

頂点 $v$ を根としたときの DP の値を返します。

計算量

$\textrm{O}(1)$

Print

void Print() const

デバッグ用に DP テーブルの内容を標準エラー出力に出力します。

計算量

$\textrm{O}(N)$

Depends on

Verified with

verify/LC-TreePathCompositeSum.test.cpp

Code

#include "Tree.hpp"

template<typename WeightType, typename Monoid>
class RerootingDP{
    public:
    using F = function<Monoid(Monoid, Monoid, Vertex)>;
    using G = function<Monoid(Monoid, WeightType, Vertex)>;
    using H = function<Monoid(Monoid, Vertex)>;
    using Fsub = function<Monoid(Monoid, Monoid)>;
    using Gsub = function<Monoid(Monoid, WeightType)>;

    RerootingDP(Graph<WeightType> &tree, Fsub merge, Gsub add, const Monoid monoid_identity, Vertex r = 0) :
            T(tree), n(tree.VertexSize()), parent(CalculateTreeParent(tree, r)), cost(CalculateTreeCost(tree, r)), child(RootedTreeAdjacentList(tree, r)),
            merge_sub_(merge), add_sub_(add), id_(monoid_identity){
        merge_ = [&](Monoid x, Monoid y, Vertex i){return merge_sub_(x, y);};
        add_ = [&](Monoid x, WeightType y, Vertex i){return add_sub_(x, y);};
        finalize_ = [](Monoid x, Vertex i){return x;};
        solve(r);
    }

    RerootingDP(Graph<WeightType> &tree, F merge, G add, H finalize, const Monoid monoid_identity, Vertex r = 0) :
            T(tree), n(tree.VertexSize()), parent(CalculateTreeParent(tree, r)), cost(CalculateTreeCost(tree, r)), child(RootedTreeAdjacentList(tree, r)),
            merge_(merge), add_(add), finalize_(finalize), id_(monoid_identity){
        solve(r);
    }

    vector<Monoid> &GetAllAnswer(){
        return dp_;
    }

    Monoid operator[](Vertex v){
        return dp_[v];
    }

    const Monoid operator[](Vertex v) const {
        return dp_[v];
    }

    void Print() const {
        cerr << "# dp table :";
        for(int i = 0; i < n; ++i){
            cerr << " " << dp_[i];
        }
        cerr << '\n';
        cerr << "# subtree_dp table" << '\n';
        for(int i = 0; i < n; ++i){
            cerr << "# vertex " << i << '\n';
            cerr << "#    subtree_dp :";
            for(int j = 0; j < subtree_dp_[i].size(); ++j){
                cerr << " " << subtree_dp_[i][j];
            }
            cerr << '\n';
            cerr << "#    left_cum   :";
            for(int j = 0; j < left_cum_[i].size(); ++j){
                cerr << " " << left_cum_[i][j];
            }
            cerr << '\n';
            cerr << "#    right_cum  :";
            for(int j = 0; j < right_cum_[i].size(); ++j){
                cerr << " " << right_cum_[i][j];
            }
            cerr << '\n';
        }
    }

    private:
    Graph<WeightType> &T;
    int n;
    vector<Vertex> parent;
    vector<WeightType> cost;
    vector<vector<Vertex>> child;
    
    vector<Monoid> dp_;
    vector<vector<Monoid>> subtree_dp_, left_cum_, right_cum_;

    const Monoid id_;

    F merge_;
    G add_;
    H finalize_;
    const Fsub merge_sub_;
    const Gsub add_sub_;

    Monoid dfs(Vertex v, bool root = false){
        Monoid ret = id_;
        for(auto u : child[v]){
            Monoid res = dfs(u);
            subtree_dp_[v].push_back(res);
            ret = merge_(ret, res, v);
        }
        if(root) ret = finalize_(ret, v);
        else ret = add_(ret, cost[v], v);
        return ret;
    }

    void solve(Vertex r){
        dp_.resize(n, id_);
        subtree_dp_.resize(n, vector<Monoid>{id_});
        left_cum_.resize(n);
        right_cum_.resize(n);

        dp_[r] = dfs(r, true);
        int root_size = subtree_dp_[r].size();
        left_cum_[r].resize(root_size + 1);
        left_cum_[r].front() = id_;
        for(int i = 1; i < root_size; ++i){
            left_cum_[r][i] = merge_(left_cum_[r][i - 1], subtree_dp_[r][i], r);
        }
        right_cum_[r].resize(root_size + 1);
        right_cum_[r].back() = id_;
        for(int i = root_size - 1; i - 1 >= 0; --i){
            right_cum_[r][i] = merge_(right_cum_[r][i + 1], subtree_dp_[r][i], r);
        }

        queue<tuple<int, int, int>> que;
        for(int i = 0; i < child[r].size(); ++i){
            que.push({child[r][i], r, i + 1});
        }
        while(que.size()){
            auto [v, p, idx] = que.front(); que.pop();
            Monoid ret = id_;
            ret = merge_(ret, left_cum_[p][idx - 1], v);
            ret = merge_(ret, right_cum_[p][idx + 1], v);
            ret = add_(ret, cost[v], p);
            subtree_dp_[v].push_back(ret);
            for(int i = 1; i + 1 < subtree_dp_[v].size(); ++i){
                ret = merge_(ret, subtree_dp_[v][i], v);
            }
            dp_[v] = finalize_(ret, v);
            int c = subtree_dp_[v].size();
            left_cum_[v].resize(c + 1);
            left_cum_[v][0] = id_;
            for(int i = 1; i < c; ++i){
                left_cum_[v][i] = merge_(left_cum_[v][i - 1], subtree_dp_[v][i], v);
            }
            right_cum_[v].resize(c + 1);
            right_cum_[v].back() = id_;
            for(int i = c - 1; i - 1 >= 0; --i){
                right_cum_[v][i] = merge_(right_cum_[v][i + 1], subtree_dp_[v][i], v);
            }
            for(int i = 0; i < child[v].size(); ++i){
                que.push({child[v][i], v, i + 1});
            }
        }
    }
};

#line 2 "Library/Tree/Tree.hpp"

#line 2 "Library/Graph/Graph.hpp"

#line 2 "Library/Common.hpp"

/**
 * @file Common.hpp
 */

#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <cstdint>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
using namespace std;

using ll = int64_t;
using ull = uint64_t;

constexpr const ll INF = (1LL << 62) - (3LL << 30) - 1;
#line 4 "Library/Graph/Graph.hpp"

using Vertex = int;

template<typename WeightType = int32_t>
struct Edge{
    public:
    Edge() = default;

    Edge(Vertex from_, Vertex to_, WeightType weight_ = 1, int idx_ = -1) :
        from(from_), to(to_), cost(weight_), idx(idx_){}
    
    bool operator<(const Edge<WeightType> &e) const {return cost < e.cost;}

    operator int() const {return to;}

    Vertex from, to;
    WeightType cost;
    int idx;
};

template<typename WeightType = int32_t>
class Graph{
    public:
    Graph() = default;

    Graph(int V) : edge_size_(0), adjacent_list_(V){}
    
    inline void AddUndirectedEdge(Vertex u, Vertex v, WeightType w = 1){
        int idx = edge_size_++;
        adjacent_list_[u].push_back(Edge<WeightType>(u, v, w, idx));
        adjacent_list_[v].push_back(Edge<WeightType>(v, u, w, idx));
    }
    
    inline void AddDirectedEdge(Vertex u, Vertex v, WeightType w = 1){
        int idx = edge_size_++;
        adjacent_list_[u].push_back(Edge<WeightType>(u, v, w, idx));
    }

    inline size_t VertexSize() const {
        return adjacent_list_.size();
    }

    inline size_t EdgeSize() const {
        return edge_size_;
    }

    inline vector<Edge<WeightType>> &operator[](const Vertex v){
        return adjacent_list_[v];
    }

    inline const vector<Edge<WeightType>> &operator[](const Vertex v) const {
        return adjacent_list_[v];
    }
    
    private:
    size_t edge_size_;
    vector<vector<Edge<WeightType>>> adjacent_list_;
};

template<typename WeightType = int32_t>
Graph<WeightType> InputGraph(int N, int M, int padding = -1, bool weighted = false, bool directed = false){
    Graph<WeightType> G(N);
    for(int i = 0; i < M; ++i){
        Vertex u, v; WeightType w = 1;
        cin >> u >> v, u += padding, v += padding;
        if(weighted) cin >> w;
        if(directed) G.AddDirectedEdge(u, v, w);
        else G.AddUndirectedEdge(u, v, w);
    }
    return G;
}
#line 4 "Library/Tree/Tree.hpp"

template<typename WeightType = int32_t>
Graph<WeightType> InputTree(int V, int padding = -1, bool weighted = false){
    Graph<WeightType> G(V);
    for(int i = 0; i < V - 1; ++i){
        Vertex u, v; WeightType w = 1;
        cin >> u >> v, u += padding, v += padding;
        if(weighted) cin >> w;
        G.AddUndirectedEdge(u, v, w);
    }
    return G;
}

template<typename WeightType = int32_t>
Graph<WeightType> InputRootedTreeChild(int V, int padding = -1){
    Graph<WeightType> G(V);
    for(Vertex u = 0; u < V; ++u){
        int k; cin >> k;
        for(int i = 0; i < k; ++i){
            Vertex v; cin >> v, v += padding;
            G.AddUndirectedEdge(u, v);
        }
    }
    return G;
}

template<typename WeightType = int32_t>
Graph<WeightType> InputRootedTreeParent(int V, int padding = -1){
    Graph<WeightType> G(V);
    for(Vertex u = 1; u < V; ++u){
        Vertex v; cin >> v, v += padding;
        G.AddUndirectedEdge(u, v);
    }
    return G;
}

template<typename WeightType = int32_t>
vector<vector<Vertex>> RootedTreeAdjacentList(const Graph<WeightType> &T, const Vertex r = 0){
    int V = T.VertexSize();
    vector<vector<Vertex>> ret(V);
    auto rec = [&](auto &self, Vertex u, Vertex p) -> void {
        for(Vertex v : T[u]){
            if(v == p) continue;
            ret[u].push_back(v);
            self(self, v, u);
        }
    };
    rec(rec, r, -1);
    return ret;
}

template<typename WeightType>
vector<Vertex> CalculateTreeParent(Graph<WeightType> &T, Vertex r = 0){
    int V = T.VertexSize();
    vector<Vertex> ret(V, -1);
    auto rec = [&](auto &self, Vertex u) -> void {
        for(Vertex v : T[u]){
            if(v == ret[u]) continue;
            ret[v] = u;
            self(self, v);
        }
    };
    rec(rec, r);
    return ret;
}

template<typename WeightType>
vector<WeightType> CalculateTreeCost(Graph<WeightType> &T, Vertex r = 0){
    int V = T.VertexSize();
    vector<WeightType> ret(V);
    auto rec = [&](auto &self, Vertex u, Vertex p) -> void {
        for(const Edge<WeightType> &e : T[u]){
            Vertex v = e.to;
            if(v == p) continue;
            ret[v] = e.cost;
            self(self, v, u);
        }
    };
    rec(rec, r, -1);
    return ret;
}

template<typename WeightType>
vector<int> CalculateTreeDepth(Graph<WeightType> &T, Vertex r = 0){
    int V = T.VertexSize();
    vector<int> ret(V, 0);
    auto rec = [&](auto &self, Vertex u, Vertex p, int d) -> void {
        ret[u] = d;
        for(Vertex v : T[u]){
            if(v == p) continue;
            self(self, v, u, d + 1);
        }
    };
    rec(rec, r, -1, 0);
    return ret;
}

template<typename WeightType>
vector<WeightType> CalculateTreeDistance(Graph<WeightType> &T, Vertex r = 0){
    int V = T.VertexSize();
    vector<WeightType> ret(V, WeightType(INF));
    auto rec = [&](auto &self, Vertex u) -> void {
        for(const Edge<WeightType> &e : T[u]){
            if(ret[e.to] > ret[u] + e.cost){
                ret[e.to] = ret[u] + e.cost;
                self(self, e.to);
            }
        }
    };
    ret[r] = 0;
    rec(rec, r);
    return ret;
}

template<typename WeightType>
vector<int> CalculateSubtreeSize(Graph<WeightType> &tree, Vertex r = 0){
    int V = tree.VertexSize();
    vector<int> ret(V, 1);
    auto rec = [&](auto self, Vertex u, Vertex p) -> int {
        for(const int v : tree[u]){
            if(v == p) continue;
            ret[u] += self(self, v, u);
        }
        return ret[u];
    };
    rec(rec, r, -1);
    return ret;
}
#line 2 "Library/Tree/RerootingDP.hpp"

template<typename WeightType, typename Monoid>
class RerootingDP{
    public:
    using F = function<Monoid(Monoid, Monoid, Vertex)>;
    using G = function<Monoid(Monoid, WeightType, Vertex)>;
    using H = function<Monoid(Monoid, Vertex)>;
    using Fsub = function<Monoid(Monoid, Monoid)>;
    using Gsub = function<Monoid(Monoid, WeightType)>;

    RerootingDP(Graph<WeightType> &tree, Fsub merge, Gsub add, const Monoid monoid_identity, Vertex r = 0) :
            T(tree), n(tree.VertexSize()), parent(CalculateTreeParent(tree, r)), cost(CalculateTreeCost(tree, r)), child(RootedTreeAdjacentList(tree, r)),
            merge_sub_(merge), add_sub_(add), id_(monoid_identity){
        merge_ = [&](Monoid x, Monoid y, Vertex i){return merge_sub_(x, y);};
        add_ = [&](Monoid x, WeightType y, Vertex i){return add_sub_(x, y);};
        finalize_ = [](Monoid x, Vertex i){return x;};
        solve(r);
    }

    RerootingDP(Graph<WeightType> &tree, F merge, G add, H finalize, const Monoid monoid_identity, Vertex r = 0) :
            T(tree), n(tree.VertexSize()), parent(CalculateTreeParent(tree, r)), cost(CalculateTreeCost(tree, r)), child(RootedTreeAdjacentList(tree, r)),
            merge_(merge), add_(add), finalize_(finalize), id_(monoid_identity){
        solve(r);
    }

    vector<Monoid> &GetAllAnswer(){
        return dp_;
    }

    Monoid operator[](Vertex v){
        return dp_[v];
    }

    const Monoid operator[](Vertex v) const {
        return dp_[v];
    }

    void Print() const {
        cerr << "# dp table :";
        for(int i = 0; i < n; ++i){
            cerr << " " << dp_[i];
        }
        cerr << '\n';
        cerr << "# subtree_dp table" << '\n';
        for(int i = 0; i < n; ++i){
            cerr << "# vertex " << i << '\n';
            cerr << "#    subtree_dp :";
            for(int j = 0; j < subtree_dp_[i].size(); ++j){
                cerr << " " << subtree_dp_[i][j];
            }
            cerr << '\n';
            cerr << "#    left_cum   :";
            for(int j = 0; j < left_cum_[i].size(); ++j){
                cerr << " " << left_cum_[i][j];
            }
            cerr << '\n';
            cerr << "#    right_cum  :";
            for(int j = 0; j < right_cum_[i].size(); ++j){
                cerr << " " << right_cum_[i][j];
            }
            cerr << '\n';
        }
    }

    private:
    Graph<WeightType> &T;
    int n;
    vector<Vertex> parent;
    vector<WeightType> cost;
    vector<vector<Vertex>> child;
    
    vector<Monoid> dp_;
    vector<vector<Monoid>> subtree_dp_, left_cum_, right_cum_;

    const Monoid id_;

    F merge_;
    G add_;
    H finalize_;
    const Fsub merge_sub_;
    const Gsub add_sub_;

    Monoid dfs(Vertex v, bool root = false){
        Monoid ret = id_;
        for(auto u : child[v]){
            Monoid res = dfs(u);
            subtree_dp_[v].push_back(res);
            ret = merge_(ret, res, v);
        }
        if(root) ret = finalize_(ret, v);
        else ret = add_(ret, cost[v], v);
        return ret;
    }

    void solve(Vertex r){
        dp_.resize(n, id_);
        subtree_dp_.resize(n, vector<Monoid>{id_});
        left_cum_.resize(n);
        right_cum_.resize(n);

        dp_[r] = dfs(r, true);
        int root_size = subtree_dp_[r].size();
        left_cum_[r].resize(root_size + 1);
        left_cum_[r].front() = id_;
        for(int i = 1; i < root_size; ++i){
            left_cum_[r][i] = merge_(left_cum_[r][i - 1], subtree_dp_[r][i], r);
        }
        right_cum_[r].resize(root_size + 1);
        right_cum_[r].back() = id_;
        for(int i = root_size - 1; i - 1 >= 0; --i){
            right_cum_[r][i] = merge_(right_cum_[r][i + 1], subtree_dp_[r][i], r);
        }

        queue<tuple<int, int, int>> que;
        for(int i = 0; i < child[r].size(); ++i){
            que.push({child[r][i], r, i + 1});
        }
        while(que.size()){
            auto [v, p, idx] = que.front(); que.pop();
            Monoid ret = id_;
            ret = merge_(ret, left_cum_[p][idx - 1], v);
            ret = merge_(ret, right_cum_[p][idx + 1], v);
            ret = add_(ret, cost[v], p);
            subtree_dp_[v].push_back(ret);
            for(int i = 1; i + 1 < subtree_dp_[v].size(); ++i){
                ret = merge_(ret, subtree_dp_[v][i], v);
            }
            dp_[v] = finalize_(ret, v);
            int c = subtree_dp_[v].size();
            left_cum_[v].resize(c + 1);
            left_cum_[v][0] = id_;
            for(int i = 1; i < c; ++i){
                left_cum_[v][i] = merge_(left_cum_[v][i - 1], subtree_dp_[v][i], v);
            }
            right_cum_[v].resize(c + 1);
            right_cum_[v].back() = id_;
            for(int i = c - 1; i - 1 >= 0; --i){
                right_cum_[v][i] = merge_(right_cum_[v][i + 1], subtree_dp_[v][i], v);
            }
            for(int i = 0; i < child[v].size(); ++i){
                que.push({child[v][i], v, i + 1});
            }
        }
    }
};

Rerooting DP - 全方位木DP (Library/Tree/RerootingDP.hpp)

Rerooting DP - 全方位木DP

Function

Constructor (辺属性のみ)

Constructor (頂点属性を含む)

GetAllAnswer

operator[]

Print

Depends on

Verified with

Code

Rerooting DP - 全方位木DP
(Library/Tree/RerootingDP.hpp)