木DP/全方位木DP
(Src/Graph/Tree/Rerooting.hpp)

• View this file on GitHub
• Last update: 2025-12-01 12:59:00+09:00
• Include: #include "Src/Graph/Tree/Rerooting.hpp"

概要

流石に木DPは自分で書いた方が早いので、用途は主に全方位木dpになると思う。

• 全方位木DPも自分で書いた方がマシかもしれないけど

テンプレートの雛形

struct DP {
    using Tree = ;
    using Cluster = ;
    using Edge = ;
    using Vertex = ;
    static Tree Convert(Vertex) {
    }
    static Cluster AddEdge(Tree,Edge) {
    }
    static Cluster Merge(Cluster,Cluster) {
    }
    static Tree AddVertex(Cluster,Vertex) {
    }
};

木dpでは

    for (int x : g[v])
        if (x != p)
            chs.push_back(rec(rec, x));
    dp[v] = prod(chs);

みたいにやると思うんですけど、それとこの雛形がどう関係しているのかというと

• Vertex: dpの計算に関する頂点の情報
• Edge: dpの計算に関する辺が持つ情報
• Cluster: prod(chs)のこと。ここでは木の根はvirtualなものと考える
• Tree: prod(chs)に頂点vの情報をマージして、根をvirtualで無いようにしたもの

みたいな感じ。

任意の木dpで全部を毎回使うわけではなくて、例えば最遠点を求める場合はVertexはdpの計算に何もかかわってこない。こういうときは適当にゴミをいれておけばいいと思います。SOLID のI(だっけ?)違反なライブラリだけど、使う側が頑張る形でいきましょう

ライブラリの使い方

template <concepts::TreeDP T>
using TreeDPGraph = std::vector<std::vector<std::pair<usize, typename T::Edge>>>;

template <concepts::TreeDP T>
std::vector<typename T::Tree> TreeDP(const TreeDPGraph<T>& g, const std::vector<typename T::Vertex>& vs, usize root) {

上の雛形にあるエイリアスとメンバが全部そろっているとcocepts::TreeDPの要求は満たされる。

gはいつもの隣接リストで、firstに行先の頂点番号、secondにその辺の情報Edgeを納めます。必ず $u\rightarrow v, v\rightarrow u$ 両方入れること。

vsは頂点の情報の列

Depends on

• 標準データ型のエイリアス (Src/Template/TypeAlias.hpp)

Verified with

• Test/AOJ/1595.test.cpp

Code

#pragma once

#include "../../Template/TypeAlias.hpp"

#include <vector>
#include <cassert>
#include <concepts>
#include <limits>

namespace zawa {

namespace concepts {

template <class T>
concept TreeDP = requires {
    typename T::Tree;
    typename T::Cluster;
    typename T::Edge;
    typename T::Vertex;
    { T::Convert(std::declval<typename T::Vertex>()) } -> std::same_as<typename T::Tree>;
    { T::AddEdge(std::declval<typename T::Tree>(), std::declval<typename T::Edge>()) } -> std::same_as<typename T::Cluster>;
    { T::Merge(std::declval<typename T::Cluster>(), std::declval<typename T::Cluster>()) } -> std::same_as<typename T::Cluster>;
    { T::AddVertex(std::declval<typename T::Cluster>(), std::declval<typename T::Vertex>()) } -> std::same_as<typename T::Tree>;
};

}

template <concepts::TreeDP T>
using TreeDPGraph = std::vector<std::vector<std::pair<usize, typename T::Edge>>>;

template <concepts::TreeDP T>
std::vector<typename T::Tree> TreeDP(const TreeDPGraph<T>& g, const std::vector<typename T::Vertex>& vs, usize root) {
    using Tree = typename T::Tree;
    using Cluster = typename T::Cluster;
    // using Edge = typename T::Edge;
    // using Vertex = typename T::Vertex;
    assert(g.size() == vs.size());
    const usize n = g.size();
    std::vector<Tree> res(n);
    auto rec = [&](auto rec, usize v, usize p) -> Tree {
        if (ssize(g[v]) == 1 and g[v][0].first == p)
            return res[v] = T::Convert(vs[v]);
        usize idx = g[v][0].first == p ? 1 : 0;
        assert(idx < g[v].size());
        Cluster cluster = T::AddEdge(rec(rec, g[v][idx].first, v), g[v][idx].second);
        idx++;
        for ( ; idx < g[v].size() ; idx++) {
            Cluster cur = T::AddEdge(rec(rec, g[v][idx].first, v), g[v][idx].second);
            cluster = T::Merge(cluster, cur);
        }
        return res[v] = T::AddVertex(cluster, vs[v]);
    };
    rec(rec, root, g.size());
    return res;
}

template <concepts::TreeDP T>
std::vector<typename T::Tree> Rerooting(const TreeDPGraph<T>& g, const std::vector<typename T::Vertex>& vs) {
    using Tree = typename T::Tree;
    using Cluster = typename T::Cluster;
    // using Edge = typename T::Edge;
    // using Vertex = typename T::Vertex;
    assert(g.size() == vs.size());
    const usize n = g.size();
    if (n <= 2) {
        std::vector<Tree> res(n);
        for (usize i = 0 ; i < n ; i++)
            res[i] = TreeDP<T>(g, vs, i)[i];
        return res;
    }
    std::vector<Cluster> dp(n);
    auto rec1 = [&](auto rec, usize v, usize p) -> Tree {
        if (ssize(g[v]) == 1 and g[v][0].first == p)
            return T::Convert(vs[v]);
        usize idx = g[v][0].first == p ? 1 : 0;
        assert(idx < g[v].size());
        dp[v] = T::AddEdge(rec(rec, g[v][idx].first, v), g[v][idx].second);
        idx++;
        for ( ; idx < g[v].size() ; idx++)
            if (g[v][idx].first != p) {
                Cluster cur = T::AddEdge(rec(rec, g[v][idx].first, v), g[v][idx].second);
                dp[v] = T::Merge(dp[v], cur);
            }
        return T::AddVertex(dp[v], vs[v]);
    };
    usize root = 0;
    while (root < g.size() and g[root].size() <= 1)
        root++;
    assert(root < g.size());
    std::vector<Tree> res(n);
    res[root] = rec1(rec1, root, g.size());
    auto rec2 = [&](auto rec, usize v, usize p, Cluster pv) -> void {
        if (ssize(g[v]) == 1) {
            assert(g[v][0].first == p);
            res[v] = T::AddVertex(pv, vs[v]);
            return;
        }
        assert(ssize(g[v]) >= 2);
        std::vector<Cluster> pref(g[v].size()), suf(g[v].size());
        pref[0] = g[v][0].first == p ? pv : T::AddEdge(dp[g[v][0].first], g[v][0].second);
        for (usize i = 1 ; i < g[v].size() ; i++) {
            Cluster cur = g[v][i].first == p ? pv : T::AddEdge(dp[g[v][i].first], g[v][i].second);
            pref[i] = T::Merge(pref[i - 1], cur);
        }
        suf[g[v].size() - 1] = g[v].back().first == p ? pv : T::AddEdge(dp[g[v].back().first], g[v].back().second);
        for (usize i = g[v].size() - 1 ; i-- ; ) {
            Cluster cur = g[v][i].first == p ? pv : T::AddEdge(dp[g[v][i].first], g[v][i].second);
            suf[i] = T::Merge(cur, suf[i + 1]);
        }
        res[v] = T::AddVertex(pref.back(), vs[v]);
        for (usize i = 0 ; i < g[v].size() ; i++)
            if (g[v][i].first != p) {
                Cluster pc = i == 0 ? suf[1] : (i + 1 == g[v].size() ? pref[i - 1] : T::Merge(pref[i - 1], suf[i + 1]));
                rec(rec, g[v][i].first, v, T::AddEdge(T::AddVertex(pc, vs[v]), g[v][i].second));
            }
    };
    rec2(rec2, root, g.size(), Cluster{});
    return res;
}

} // namespace zawa

#line 2 "Src/Graph/Tree/Rerooting.hpp"

#line 2 "Src/Template/TypeAlias.hpp"

#include <cstdint>
#include <cstddef>

namespace zawa {

using i16 = std::int16_t;
using i32 = std::int32_t;
using i64 = std::int64_t;
using i128 = __int128_t;

using u8 = std::uint8_t;
using u16 = std::uint16_t;
using u32 = std::uint32_t;
using u64 = std::uint64_t;

using usize = std::size_t;

} // namespace zawa
#line 4 "Src/Graph/Tree/Rerooting.hpp"

#include <vector>
#include <cassert>
#include <concepts>
#include <limits>

namespace zawa {

namespace concepts {

template <class T>
concept TreeDP = requires {
    typename T::Tree;
    typename T::Cluster;
    typename T::Edge;
    typename T::Vertex;
    { T::Convert(std::declval<typename T::Vertex>()) } -> std::same_as<typename T::Tree>;
    { T::AddEdge(std::declval<typename T::Tree>(), std::declval<typename T::Edge>()) } -> std::same_as<typename T::Cluster>;
    { T::Merge(std::declval<typename T::Cluster>(), std::declval<typename T::Cluster>()) } -> std::same_as<typename T::Cluster>;
    { T::AddVertex(std::declval<typename T::Cluster>(), std::declval<typename T::Vertex>()) } -> std::same_as<typename T::Tree>;
};

}

template <concepts::TreeDP T>
using TreeDPGraph = std::vector<std::vector<std::pair<usize, typename T::Edge>>>;

template <concepts::TreeDP T>
std::vector<typename T::Tree> TreeDP(const TreeDPGraph<T>& g, const std::vector<typename T::Vertex>& vs, usize root) {
    using Tree = typename T::Tree;
    using Cluster = typename T::Cluster;
    // using Edge = typename T::Edge;
    // using Vertex = typename T::Vertex;
    assert(g.size() == vs.size());
    const usize n = g.size();
    std::vector<Tree> res(n);
    auto rec = [&](auto rec, usize v, usize p) -> Tree {
        if (ssize(g[v]) == 1 and g[v][0].first == p)
            return res[v] = T::Convert(vs[v]);
        usize idx = g[v][0].first == p ? 1 : 0;
        assert(idx < g[v].size());
        Cluster cluster = T::AddEdge(rec(rec, g[v][idx].first, v), g[v][idx].second);
        idx++;
        for ( ; idx < g[v].size() ; idx++) {
            Cluster cur = T::AddEdge(rec(rec, g[v][idx].first, v), g[v][idx].second);
            cluster = T::Merge(cluster, cur);
        }
        return res[v] = T::AddVertex(cluster, vs[v]);
    };
    rec(rec, root, g.size());
    return res;
}

template <concepts::TreeDP T>
std::vector<typename T::Tree> Rerooting(const TreeDPGraph<T>& g, const std::vector<typename T::Vertex>& vs) {
    using Tree = typename T::Tree;
    using Cluster = typename T::Cluster;
    // using Edge = typename T::Edge;
    // using Vertex = typename T::Vertex;
    assert(g.size() == vs.size());
    const usize n = g.size();
    if (n <= 2) {
        std::vector<Tree> res(n);
        for (usize i = 0 ; i < n ; i++)
            res[i] = TreeDP<T>(g, vs, i)[i];
        return res;
    }
    std::vector<Cluster> dp(n);
    auto rec1 = [&](auto rec, usize v, usize p) -> Tree {
        if (ssize(g[v]) == 1 and g[v][0].first == p)
            return T::Convert(vs[v]);
        usize idx = g[v][0].first == p ? 1 : 0;
        assert(idx < g[v].size());
        dp[v] = T::AddEdge(rec(rec, g[v][idx].first, v), g[v][idx].second);
        idx++;
        for ( ; idx < g[v].size() ; idx++)
            if (g[v][idx].first != p) {
                Cluster cur = T::AddEdge(rec(rec, g[v][idx].first, v), g[v][idx].second);
                dp[v] = T::Merge(dp[v], cur);
            }
        return T::AddVertex(dp[v], vs[v]);
    };
    usize root = 0;
    while (root < g.size() and g[root].size() <= 1)
        root++;
    assert(root < g.size());
    std::vector<Tree> res(n);
    res[root] = rec1(rec1, root, g.size());
    auto rec2 = [&](auto rec, usize v, usize p, Cluster pv) -> void {
        if (ssize(g[v]) == 1) {
            assert(g[v][0].first == p);
            res[v] = T::AddVertex(pv, vs[v]);
            return;
        }
        assert(ssize(g[v]) >= 2);
        std::vector<Cluster> pref(g[v].size()), suf(g[v].size());
        pref[0] = g[v][0].first == p ? pv : T::AddEdge(dp[g[v][0].first], g[v][0].second);
        for (usize i = 1 ; i < g[v].size() ; i++) {
            Cluster cur = g[v][i].first == p ? pv : T::AddEdge(dp[g[v][i].first], g[v][i].second);
            pref[i] = T::Merge(pref[i - 1], cur);
        }
        suf[g[v].size() - 1] = g[v].back().first == p ? pv : T::AddEdge(dp[g[v].back().first], g[v].back().second);
        for (usize i = g[v].size() - 1 ; i-- ; ) {
            Cluster cur = g[v][i].first == p ? pv : T::AddEdge(dp[g[v][i].first], g[v][i].second);
            suf[i] = T::Merge(cur, suf[i + 1]);
        }
        res[v] = T::AddVertex(pref.back(), vs[v]);
        for (usize i = 0 ; i < g[v].size() ; i++)
            if (g[v][i].first != p) {
                Cluster pc = i == 0 ? suf[1] : (i + 1 == g[v].size() ? pref[i - 1] : T::Merge(pref[i - 1], suf[i + 1]));
                rec(rec, g[v][i].first, v, T::AddEdge(T::AddVertex(pc, vs[v]), g[v][i].second));
            }
    };
    rec2(rec2, root, g.size(), Cluster{});
    return res;
}

} // namespace zawa

Back to top page