重軽再帰 DP (graph/tree/hl-rec-dp.hpp)

• View this file on GitHub
• View document part on GitHub
• Last update: 2026-06-07 20:36:35+09:00
• Include: #include "graph/tree/hl-rec-dp.hpp"
• Link: https://arxiv.org/abs/1807.04942

概要

• 根付き木の部分木を、1 本の State から K 本の State の組へ写す変換として再帰的に計算する。
• 各頂点では最大部分木サイズの子を重い子とし、重い子だけは全レーンをまとめて 1 回再帰する。
• 軽い子は各レーンごとに再帰し、Spec 側の関数で親側の状態に戻す。
• before_vertex と after_vertex により、固定した根のもとで各頂点を部分木の根とする値を列挙できる。

使い方

• hl_rec_dp(n, edges, root, initial_state, spec)
  • 頂点数が n で、無向辺列が edges である木を root で根付き木にして実行する。
  • 頂点番号は 0-based indexing とする。
  • edges はサイズ $n-1$ の木の辺列である。
  • initial_state は各重パスの根側から渡す初期状態である。
  • spec の before_vertex, add_vertex, after_vertex などを呼び出し、spec の中の答えの配列を更新する。
  • 返り値は root を通常の頂点として加えた後の std::array<Spec::State, Spec::K> である。
  • 前提：$n\ge 1$、$0\le \mathrm{root}<n$、edges は木である。
• Spec
  • using State = ...; と static constexpr int K = ...; を持つ。
  • make_pack(v, in) は、子をまだ処理していない頂点 v の DP の組を返す。
  • take_heavy(v, child, pack) は、重い子 child の返り値を頂点 v 側の DP の組に変換して返す。
  • take_light(v, child, lane, pack) は、軽い子 child をレーン lane だけ処理した結果から、親側の 1 本の State を返す。
  • before_vertex(v, pack) は、頂点 v を部分木根として扱う値を記録するために呼ばれる。
  • add_vertex(v, pack) は、頂点 v を親へ返す通常の頂点として pack に反映する。
  • after_vertex(v, pack) は、add_vertex 後の値を必要に応じて記録するために呼ばれる。
  • pack の型は std::array<State, K> である。

使用例（問題のネタバレを含む）

例として、AtCoder Beginner Contest 311 Ex - Many Illumination Plans を解くソースコードを示す。

#include <bits/stdc++.h>

#include "graph/tree/hl-rec-dp.hpp"

using namespace std;
using ll = long long;

#define rep(i, r) for (int i = 0; i < (int)(r); ++i)

template <class T1, class T2> bool chmax(T1 &l, const T2 &r) {
    if (r > l) {
        l = r;
        return true;
    }
    return false;
}

namespace abc311 {
struct Spec {
    static constexpr int K = 2;
    using State = vector<ll>;
    using Pack = array<State, K>;

    static constexpr ll NEG = -2'000'000'000'000'000'001;

    vector<ll> ans;

    int x = 0;
    const vector<ll> &b;
    const vector<int> &w, c;

    Spec(int n, int x, const vector<ll> &b, const vector<int> &w,
         const vector<int> &c)
        : x(x), b(b), w(w), c(c) {
        ans.assign(n, NEG);
    }

    Pack make_pack(int, const State &in) { return Pack{in, in}; }

    Pack take_heavy(int, int, Pack &&child_dp) { return move(child_dp); }

    State take_light(int, int, int lane, Pack &&child_dp) {
        return move(child_dp[lane]);
    }

    void before_vertex(int v, const Pack &dp) {
        int c_v = c[v];
        int w_v = w[v];
        ll best = NEG;
        rep(i, x - w_v + 1) { chmax(best, dp[c_v ^ 1][i] + b[v]); }
        chmax(ans[v], best);
    }

    void add_vertex(int v, Pack &dp) {
        int c_v = c[v];
        int w_v = w[v];
        rep(i, x - w_v + 1) { chmax(dp[c_v][i + w_v], dp[c_v ^ 1][i] + b[v]); }
    }

    void after_vertex(int, const Pack &) {}
};
} // namespace abc311

int main() {
    int n, x;
    cin >> n >> x;

    vector<pair<int, int>> edges(n - 1);
    rep(i, n - 1) {
        int p_i;
        cin >> p_i;
        p_i -= 1;
        edges[i] = {i + 1, p_i};
    }

    vector<int> w(n), c(n);
    vector<ll> b(n);
    rep(i, n) { cin >> b[i] >> w[i] >> c[i]; }

    using abc311::Spec;

    Spec::State initial_state(x + 1, Spec::NEG);
    initial_state[0] = 0;
    Spec spec(n, x, b, w, c);

    hl_rec_dp(n, edges, 0, initial_state, spec);

    rep(i, n) { cout << spec.ans[i] << "\n"; }

    return 0;
}

計算量

• 木の根付き化と重い子の選択：$O(n)$。
• 軽い辺だけで再帰が深くなるため、再帰段数：$O(\log n)$。
• $K$ を Spec::K の値とする。部分木 1 回分の値計算の呼び出し回数は $O(n^{\log_2(K+1)})$。
• 全頂点の before_vertex / after_vertex を回収する実行では、$K\ge 2$ なら同じく $O(n^{\log_2(K+1)})$ 回、$K=1$ なら $O(n\log n)$ 回である。
• 各呼び出しでの Spec 側の処理時間が掛かる。典型的に $K=2$、各処理が $O(X)$ なら $O(n^{\log_2 3}X)$。

参考文献

1. Soh Kumabe, Takanori Maehara, and Ryoma Sin’ya. Linear Pseudo-Polynomial Factor Algorithm for Automaton Constrained Tree Knapsack Problem. In WALCOM: Algorithms and Computation, Lecture Notes in Computer Science, Vol. 11355, pp. 248–260. Springer, 2019. doi:10.1007/978-3-030-10564-8_20.（arXiv のリンク）

Verified with

• verify/standalone-hl-rec-dp.test.cpp

Code

#ifndef GRAPH_TREE_HL_REC_DP_HPP
#define GRAPH_TREE_HL_REC_DP_HPP

// 重軽分解の重い子だけを同じ初期状態で再帰する DP の骨格である。
// Spec は State, K と 6 個の関数を持つ必要がある。
// 入力は木、または根から全頂点へ到達する根付き木である。
// 軽い子へ潜る再帰だけを積むため、再帰の段数は O(log N) である。
// 計算量は Spec の処理時間と呼び出し回数に依存する。

// 参考文献:
// 1. Soh Kumabe, Takanori Maehara, and Ryoma Sin'ya. Linear Pseudo-Polynomial Factor Algorithm for Automaton Constrained Tree Knapsack Problem. In WALCOM: Algorithms and Computation, Lecture Notes in Computer Science, Vol. 11355, pp. 248–260. Springer, 2019. doi:10.1007/978-3-030-10564-8_20.
//    arXiv: https://arxiv.org/abs/1807.04942

#include <array>
#include <cassert>
#include <utility>
#include <vector>

namespace hl_rec_dp_impl {
inline void put_heavy_child_first(std::vector<std::vector<int>> &child,
                                  const std::vector<int> &order) {
    const int n = static_cast<int>(child.size());
    std::vector<int> subtree_size(n, 1);
    for (int i = static_cast<int>(order.size()) - 1; i >= 0; --i) {
        const int v = order[i];
        for (int to : child[v]) {
            subtree_size[v] += subtree_size[to];
        }
        if (!child[v].empty()) {
            int heavy_index = 0;
            for (int j = 1; j < static_cast<int>(child[v].size()); ++j) {
                const int a = child[v][j];
                const int b = child[v][heavy_index];
                if (subtree_size[a] > subtree_size[b] ||
                    (subtree_size[a] == subtree_size[b] && a < b)) {
                    heavy_index = j;
                }
            }
            std::swap(child[v][0], child[v][heavy_index]);
        }
    }
}

inline std::vector<int>
validate_rooted_tree(const std::vector<std::vector<int>> &child, int root) {
    const int n = static_cast<int>(child.size());
    assert(n >= 1);
    assert(0 <= root && root < n);
    std::vector<int> parent(n, -2);
    std::vector<int> order;
    order.reserve(n);
    std::vector<int> stack{root};
    parent[root] = -1;
    while (!stack.empty()) {
        const int v = stack.back();
        stack.pop_back();
        order.push_back(v);
        for (int to : child[v]) {
            assert(0 <= to && to < n);
            assert(parent[to] == -2);
            parent[to] = v;
            stack.push_back(to);
        }
    }
    assert(static_cast<int>(order.size()) == n);
    return order;
}

inline std::vector<std::vector<int>>
make_rooted_tree(int n, const std::vector<std::pair<int, int>> &edges,
                 int root) {
    assert(n >= 1);
    assert(static_cast<int>(edges.size()) == n - 1);
    assert(0 <= root && root < n);
    std::vector<std::vector<int>> graph(n);
    for (const auto &[u, v] : edges) {
        assert(0 <= u && u < n);
        assert(0 <= v && v < n);
        assert(u != v);
        graph[u].push_back(v);
        graph[v].push_back(u);
    }

    std::vector<std::vector<int>> child(n);
    std::vector<int> parent(n, -2);
    std::vector<int> order;
    order.reserve(n);
    std::vector<int> stack{root};
    parent[root] = -1;
    while (!stack.empty()) {
        const int v = stack.back();
        stack.pop_back();
        order.push_back(v);
        for (int to : graph[v]) {
            if (to == parent[v]) {
                continue;
            }
            assert(parent[to] == -2);
            parent[to] = v;
            child[v].push_back(to);
            stack.push_back(to);
        }
    }
    assert(static_cast<int>(order.size()) == n);
    put_heavy_child_first(child, order);
    return child;
}

template <class Spec> struct hl_rec_dp_runner {
    using State = typename Spec::State;
    static constexpr int K = Spec::K;
    static_assert(K >= 1);
    using Pack = std::array<State, K>;

    std::vector<std::vector<int>> child;
    Spec &spec;

    hl_rec_dp_runner(std::vector<std::vector<int>> child, Spec &spec)
        : child(std::move(child)), spec(spec) {}

    Pack run(int root, const State &initial_state) {
        return dfs_heavy_path(root, initial_state, true);
    }

    Pack dfs_heavy_path(int start, const State &in, bool collect) {
        std::vector<int> path;
        for (int v = start;; v = child[v][0]) {
            path.push_back(v);
            if (child[v].empty()) {
                break;
            }
        }

        Pack cur = spec.make_pack(path.back(), in);
        for (int i = static_cast<int>(path.size()) - 1; i >= 0; --i) {
            const int v = path[i];
            if (i + 1 < static_cast<int>(path.size())) {
                cur = spec.take_heavy(v, path[i + 1], std::move(cur));
            }
            for (int idx = 1; idx < static_cast<int>(child[v].size()); ++idx) {
                const int to = child[v][idx];
                for (int lane = 0; lane < K; ++lane) {
                    Pack got = dfs_heavy_path(to, cur[lane], false);
                    cur[lane] = spec.take_light(v, to, lane, std::move(got));
                }
            }
            if (collect) {
                spec.before_vertex(v, cur);
            }
            spec.add_vertex(v, cur);
            if (collect) {
                spec.after_vertex(v, cur);
                for (int idx = 1; idx < static_cast<int>(child[v].size());
                     ++idx) {
                    dfs_heavy_path(child[v][idx], in, true);
                }
            }
        }
        return cur;
    }
};
} // namespace hl_rec_dp_impl

template <class Spec>
std::array<typename Spec::State, Spec::K>
hl_rec_dp(int n, const std::vector<std::pair<int, int>> &edges, int root,
          const typename Spec::State &initial_state, Spec &spec) {
    auto child = hl_rec_dp_impl::make_rooted_tree(n, edges, root);
    hl_rec_dp_impl::hl_rec_dp_runner<Spec> runner(std::move(child), spec);
    return runner.run(root, initial_state);
}

#endif

#line 1 "graph/tree/hl-rec-dp.hpp"



// 重軽分解の重い子だけを同じ初期状態で再帰する DP の骨格である。
// Spec は State, K と 6 個の関数を持つ必要がある。
// 入力は木、または根から全頂点へ到達する根付き木である。
// 軽い子へ潜る再帰だけを積むため、再帰の段数は O(log N) である。
// 計算量は Spec の処理時間と呼び出し回数に依存する。

// 参考文献:
// 1. Soh Kumabe, Takanori Maehara, and Ryoma Sin'ya. Linear Pseudo-Polynomial Factor Algorithm for Automaton Constrained Tree Knapsack Problem. In WALCOM: Algorithms and Computation, Lecture Notes in Computer Science, Vol. 11355, pp. 248–260. Springer, 2019. doi:10.1007/978-3-030-10564-8_20.
//    arXiv: https://arxiv.org/abs/1807.04942

#include <array>
#include <cassert>
#include <utility>
#include <vector>

namespace hl_rec_dp_impl {
inline void put_heavy_child_first(std::vector<std::vector<int>> &child,
                                  const std::vector<int> &order) {
    const int n = static_cast<int>(child.size());
    std::vector<int> subtree_size(n, 1);
    for (int i = static_cast<int>(order.size()) - 1; i >= 0; --i) {
        const int v = order[i];
        for (int to : child[v]) {
            subtree_size[v] += subtree_size[to];
        }
        if (!child[v].empty()) {
            int heavy_index = 0;
            for (int j = 1; j < static_cast<int>(child[v].size()); ++j) {
                const int a = child[v][j];
                const int b = child[v][heavy_index];
                if (subtree_size[a] > subtree_size[b] ||
                    (subtree_size[a] == subtree_size[b] && a < b)) {
                    heavy_index = j;
                }
            }
            std::swap(child[v][0], child[v][heavy_index]);
        }
    }
}

inline std::vector<int>
validate_rooted_tree(const std::vector<std::vector<int>> &child, int root) {
    const int n = static_cast<int>(child.size());
    assert(n >= 1);
    assert(0 <= root && root < n);
    std::vector<int> parent(n, -2);
    std::vector<int> order;
    order.reserve(n);
    std::vector<int> stack{root};
    parent[root] = -1;
    while (!stack.empty()) {
        const int v = stack.back();
        stack.pop_back();
        order.push_back(v);
        for (int to : child[v]) {
            assert(0 <= to && to < n);
            assert(parent[to] == -2);
            parent[to] = v;
            stack.push_back(to);
        }
    }
    assert(static_cast<int>(order.size()) == n);
    return order;
}

inline std::vector<std::vector<int>>
make_rooted_tree(int n, const std::vector<std::pair<int, int>> &edges,
                 int root) {
    assert(n >= 1);
    assert(static_cast<int>(edges.size()) == n - 1);
    assert(0 <= root && root < n);
    std::vector<std::vector<int>> graph(n);
    for (const auto &[u, v] : edges) {
        assert(0 <= u && u < n);
        assert(0 <= v && v < n);
        assert(u != v);
        graph[u].push_back(v);
        graph[v].push_back(u);
    }

    std::vector<std::vector<int>> child(n);
    std::vector<int> parent(n, -2);
    std::vector<int> order;
    order.reserve(n);
    std::vector<int> stack{root};
    parent[root] = -1;
    while (!stack.empty()) {
        const int v = stack.back();
        stack.pop_back();
        order.push_back(v);
        for (int to : graph[v]) {
            if (to == parent[v]) {
                continue;
            }
            assert(parent[to] == -2);
            parent[to] = v;
            child[v].push_back(to);
            stack.push_back(to);
        }
    }
    assert(static_cast<int>(order.size()) == n);
    put_heavy_child_first(child, order);
    return child;
}

template <class Spec> struct hl_rec_dp_runner {
    using State = typename Spec::State;
    static constexpr int K = Spec::K;
    static_assert(K >= 1);
    using Pack = std::array<State, K>;

    std::vector<std::vector<int>> child;
    Spec &spec;

    hl_rec_dp_runner(std::vector<std::vector<int>> child, Spec &spec)
        : child(std::move(child)), spec(spec) {}

    Pack run(int root, const State &initial_state) {
        return dfs_heavy_path(root, initial_state, true);
    }

    Pack dfs_heavy_path(int start, const State &in, bool collect) {
        std::vector<int> path;
        for (int v = start;; v = child[v][0]) {
            path.push_back(v);
            if (child[v].empty()) {
                break;
            }
        }

        Pack cur = spec.make_pack(path.back(), in);
        for (int i = static_cast<int>(path.size()) - 1; i >= 0; --i) {
            const int v = path[i];
            if (i + 1 < static_cast<int>(path.size())) {
                cur = spec.take_heavy(v, path[i + 1], std::move(cur));
            }
            for (int idx = 1; idx < static_cast<int>(child[v].size()); ++idx) {
                const int to = child[v][idx];
                for (int lane = 0; lane < K; ++lane) {
                    Pack got = dfs_heavy_path(to, cur[lane], false);
                    cur[lane] = spec.take_light(v, to, lane, std::move(got));
                }
            }
            if (collect) {
                spec.before_vertex(v, cur);
            }
            spec.add_vertex(v, cur);
            if (collect) {
                spec.after_vertex(v, cur);
                for (int idx = 1; idx < static_cast<int>(child[v].size());
                     ++idx) {
                    dfs_heavy_path(child[v][idx], in, true);
                }
            }
        }
        return cur;
    }
};
} // namespace hl_rec_dp_impl

template <class Spec>
std::array<typename Spec::State, Spec::K>
hl_rec_dp(int n, const std::vector<std::pair<int, int>> &edges, int root,
          const typename Spec::State &initial_state, Spec &spec) {
    auto child = hl_rec_dp_impl::make_rooted_tree(n, edges, root);
    hl_rec_dp_impl::hl_rec_dp_runner<Spec> runner(std::move(child), spec);
    return runner.run(root, initial_state);
}

Back to top page