library/tree/ReRooting.hpp

View this file on GitHub
Last update: 2025-11-24 18:49:07+09:00
Include: #include "library/tree/ReRooting.hpp"

Depends on

build/pch/stdc++.hpp

Verified with

test/AOJ/GRL_5_A.test.cpp

Code

template <typename TREE, typename Data> class ReRooting {
    int n;
    TREE T;

  public:
    ReRooting(const TREE &T) : T(T), n(T.n) {}

    template <typename F1, typename F2>
    std::vector<Data> build(const F1 &score, const F2 &merge,
                            const Data &unit) {
        if (!~T.root)
            T.build(0);
        std::vector<Data> dp1(n, unit), dp2(n);
        for (int v : T.DFS)
            for (const auto &e : T.son(v)) {
                dp2[e.to] = score(dp1[e.to], e);
                merge(dp1[v], dp2[e.to]);
            }
        std::vector<Data> top(n, unit), res(n);
        for (int v : T.BFS) {
            Data now = (T.root == v ? unit : score(top[v], T.parent(v)));
            for (int to : T.son(v)) {
                top[to] = now;
                merge(now, dp2[to]);
            }
            res[v] = now;
            now = unit;
            for (int i = T.son(v).size() - 1; i >= 0; i--) {
                int to = T.son(v)[i];
                merge(top[to], now);
                merge(now, dp2[to]);
            }
        }
        return res;
    }
};

#line 1 "library/tree/ReRooting.hpp"
template <typename TREE, typename Data> class ReRooting {
    int n;
    TREE T;

  public:
    ReRooting(const TREE &T) : T(T), n(T.n) {}

    template <typename F1, typename F2>
    std::vector<Data> build(const F1 &score, const F2 &merge,
                            const Data &unit) {
        if (!~T.root)
            T.build(0);
        std::vector<Data> dp1(n, unit), dp2(n);
        for (int v : T.DFS)
            for (const auto &e : T.son(v)) {
                dp2[e.to] = score(dp1[e.to], e);
                merge(dp1[v], dp2[e.to]);
            }
        std::vector<Data> top(n, unit), res(n);
        for (int v : T.BFS) {
            Data now = (T.root == v ? unit : score(top[v], T.parent(v)));
            for (int to : T.son(v)) {
                top[to] = now;
                merge(now, dp2[to]);
            }
            res[v] = now;
            now = unit;
            for (int i = T.son(v).size() - 1; i >= 0; i--) {
                int to = T.son(v)[i];
                merge(top[to], now);
                merge(now, dp2[to]);
            }
        }
        return res;
    }
};