Rerooting
(src/tree/Rerooting.hpp)

• View this file on GitHub
• Last update: 2024-10-28 15:43:27+09:00
• Include: #include "src/tree/Rerooting.hpp"
• Link: View error logs on GitHub Actions

概要

全方位木 dp を計算する．

使用する際には，以下の関数をもつ構造体 TreeDP が必要となる．

• T e(): 単位元（根が virtual で辺のない 1 頂点の木）を返す．
• T op(const T& l, const T& r): 根が virtual な木 $l, r$ の合成を返す．
• T add_vertex(const T& t, int v): 根が virtual な木 $t$ の根に $v$ を代入した木を返す．
• T add_edge(const T& t, int e): 木 $t$ の根に枝 $e$ を接続してできる，根が virtual な木を返す．

struct TreeDP {
    struct T {};

    T e() {}

    T op(const T& l, const T& r) {}

    T add_vertex(const T& t, int v) {}

    T add_edge(const T& t, int e) {}
};

計算量は時間・空間計算量ともに $\mathrm{O}(n)$．

引数の STORES_SUBTREE を true にすることで，木全体から辺 $(u, v)$ を取り除いた際の $u$ を根とする部分木の情報を $get(u, v)$ で得ることができる． この場合，時間計算量は $\mathrm{O}(n \log n)$ となる．

Links

• 全方位木DP (ReRooting) の抽象化について - メモ帳

出題例

• The 2022 ICPC Asia Hangzhou Regional Programming Contest M. Please Save Pigeland
• AtCoder Regular Contest 179 D - Portable Gate

Verified with

• test/yosupo/tree_path_composite_sum.test.cpp

Code

#pragma once
#include <cassert>
#include <map>
#include <vector>

template <class TreeDP, bool STORE_SUBTREE = false> struct Rerooting {
    using T = TreeDP::T;

    struct edge {
        int to, e;
        T dp, ndp;
        edge(int to, int e, T dp, T ndp) : to(to), e(e), dp(dp), ndp(ndp) {}
    };

    std::vector<std::vector<edge>> G;
    std::unordered_map<long long, T> mp;

    Rerooting(int n, const TreeDP& treedp) : n(n), m(0), G(n), treedp(treedp), sub(n), dp(n) {}

    void add_edge(int u, int v) {
        assert(0 <= u and u < n);
        assert(0 <= v and v < n);
        G[u].emplace_back(v, m, treedp.e(), treedp.e());
        G[v].emplace_back(u, m, treedp.e(), treedp.e());
        m++;
    }

    std::vector<T> solve() {
        dfs_sub(0, -1);
        dfs_all(0, -1, treedp.e());
        return dp;
    }

    T get(int u, int v) {
        assert(STORE_SUBTREE);
        return mp[1LL * n * u + v];
    }

  private:
    int n, m;
    TreeDP treedp;
    std::vector<T> sub, dp;

    void dfs_sub(int v, int p) {
        sub[v] = treedp.e();
        for (auto& e : G[v]) {
            if (e.to == p) continue;
            dfs_sub(e.to, v);
            sub[v] = treedp.op(sub[v], treedp.add_edge(sub[e.to], e.e));
        }
        sub[v] = treedp.add_vertex(sub[v], v);
    }

    void dfs_all(int v, int p, const T& top) {
        T cur = treedp.e();
        for (int i = 0; i < int(G[v].size()); i++) {
            auto& e = G[v][i];
            e.ndp = cur;
            e.dp = treedp.add_edge(e.to == p ? top : sub[e.to], e.e);
            cur = treedp.op(cur, e.dp);
        }
        dp[v] = treedp.add_vertex(cur, v);
        cur = treedp.e();
        for (int i = int(G[v].size()) - 1; i >= 0; i--) {
            auto& e = G[v][i];
            e.ndp = treedp.add_vertex(treedp.op(e.ndp, cur), v);
            if (e.to != p) {
                if (STORE_SUBTREE) {
                    mp[1LL * n * v + e.to] = e.ndp;
                    mp[1LL * n * e.to + v] = sub[e.to];
                }
                dfs_all(e.to, v, e.ndp);
            }
            cur = treedp.op(cur, e.dp);
        }
    }
};

#line 2 "src/tree/Rerooting.hpp"
#include <cassert>
#include <map>
#include <vector>

template <class TreeDP, bool STORE_SUBTREE = false> struct Rerooting {
    using T = TreeDP::T;

    struct edge {
        int to, e;
        T dp, ndp;
        edge(int to, int e, T dp, T ndp) : to(to), e(e), dp(dp), ndp(ndp) {}
    };

    std::vector<std::vector<edge>> G;
    std::unordered_map<long long, T> mp;

    Rerooting(int n, const TreeDP& treedp) : n(n), m(0), G(n), treedp(treedp), sub(n), dp(n) {}

    void add_edge(int u, int v) {
        assert(0 <= u and u < n);
        assert(0 <= v and v < n);
        G[u].emplace_back(v, m, treedp.e(), treedp.e());
        G[v].emplace_back(u, m, treedp.e(), treedp.e());
        m++;
    }

    std::vector<T> solve() {
        dfs_sub(0, -1);
        dfs_all(0, -1, treedp.e());
        return dp;
    }

    T get(int u, int v) {
        assert(STORE_SUBTREE);
        return mp[1LL * n * u + v];
    }

  private:
    int n, m;
    TreeDP treedp;
    std::vector<T> sub, dp;

    void dfs_sub(int v, int p) {
        sub[v] = treedp.e();
        for (auto& e : G[v]) {
            if (e.to == p) continue;
            dfs_sub(e.to, v);
            sub[v] = treedp.op(sub[v], treedp.add_edge(sub[e.to], e.e));
        }
        sub[v] = treedp.add_vertex(sub[v], v);
    }

    void dfs_all(int v, int p, const T& top) {
        T cur = treedp.e();
        for (int i = 0; i < int(G[v].size()); i++) {
            auto& e = G[v][i];
            e.ndp = cur;
            e.dp = treedp.add_edge(e.to == p ? top : sub[e.to], e.e);
            cur = treedp.op(cur, e.dp);
        }
        dp[v] = treedp.add_vertex(cur, v);
        cur = treedp.e();
        for (int i = int(G[v].size()) - 1; i >= 0; i--) {
            auto& e = G[v][i];
            e.ndp = treedp.add_vertex(treedp.op(e.ndp, cur), v);
            if (e.to != p) {
                if (STORE_SUBTREE) {
                    mp[1LL * n * v + e.to] = e.ndp;
                    mp[1LL * n * e.to + v] = sub[e.to];
                }
                dfs_all(e.to, v, e.ndp);
            }
            cur = treedp.op(cur, e.dp);
        }
    }
};

Back to top page