#772882 (C++17) No.399 動的な領主

• 問題: No.399 動的な領主
• ユーザー: suisen
• 提出日時: 2022-06-23 02:54:43
• 言語: C++17 (gcc 13.3.0 + boost 1.87.0)
• 結果: AC
• 実行時間: 296 ms / 2,000 ms
• コード長: 15,243 bytes
• コンパイル時間: 1,650 ms
• コンパイル使用メモリ: 102,692 KB
• 最終ジャッジ日時: 2025-01-29 23:40:25
• ジャッジサーバーID （参考情報）: judge3 / judge5

ファイルパターン	結果
other	AC * 19

ソースコード

#define PROBLEM "https://yukicoder.me/problems/no/399"

#include <iostream>

#include <cassert>
#include <optional>
#include <utility>
#include <vector>

#include <deque>

namespace suisen {
    template <typename T, bool auto_extend = false>
    struct ObjectPool {
        using value_type = T;
        using value_pointer_type = T*;

        template <typename U>
        using container_type = std::conditional_t<auto_extend, std::deque<U>, std::vector<U>>;

        container_type<value_type> pool;
        container_type<value_pointer_type> stock;
        decltype(stock.begin()) it;

        ObjectPool() : ObjectPool(0) {}
        ObjectPool(int siz) : pool(siz), stock(siz) {
            clear();
        }

        int capacity() const { return pool.size(); }
        int size() const { return it - stock.begin(); }

        value_pointer_type alloc() {
            if constexpr (auto_extend) ensure();
            return *it++;
        }

        void free(value_pointer_type t) {
            *--it = t;
        }

        void clear() {
            int siz = pool.size();
            it = stock.begin();
            for (int i = 0; i < siz; i++) stock[i] = &pool[i];
        }

        void ensure() {
            if (it != stock.end()) return;
            int siz = stock.size();
            for (int i = siz; i <= siz * 2; ++i) {
                stock.push_back(&pool.emplace_back());
            }
            it = stock.begin() + siz;
        }
    };
} // namespace suisen

namespace suisen::internal::link_cut_tree {
    template <typename T, typename Derived>
    struct SplayTreeNodeBase {
        friend Derived;
        template <typename, typename>
        friend struct LinkCutTreeBase;

        using value_type = T;
        using node_type = Derived;
        using node_pointer_type = node_type*;

        explicit SplayTreeNodeBase(const value_type& val = value_type{}) : _val(val) {}

    protected:
        node_pointer_type _p = nullptr;
        node_pointer_type _ch[2]{ nullptr, nullptr };

        int _siz = 1;
        value_type _val;

        bool _rev = false;
        
        static bool is_root(node_pointer_type node) {
            return not node->_p or (node->_p->_ch[0] != node and node->_p->_ch[1] != node);
        }
        static node_pointer_type& parent(node_pointer_type node) {
            return node->_p;
        }
        static node_pointer_type& child(node_pointer_type node, int ch_idx) {
            return node->_ch[ch_idx];
        }
        static int size(node_pointer_type node) {
            return node ? node->_siz : 0;
        }
        static const value_type& value(node_pointer_type node) {
            return node->_val;
        }
        static void set_value(node_pointer_type node, const value_type &new_val) {
            node->_val = new_val;
        }
        static void update(node_pointer_type node) {
            node->_siz = 1 + node_type::size(node->_ch[0]) + node_type::size(node->_ch[1]);
        }

        static void reverse_all(node_pointer_type node) {
            if (not node) return;
            node->_rev ^= true;
            std::swap(node->_ch[0], node->_ch[1]);
        }

        static void push(node_pointer_type node) {
            if (std::exchange(node->_rev, false)) {
                node_type::reverse_all(node->_ch[0]);
                node_type::reverse_all(node->_ch[1]);
            }
        }

        static void rot(node_pointer_type node, int ch_idx) {
            assert(node->_ch[ch_idx]);

            node_pointer_type rt = node->_ch[ch_idx];
            if (not node_type::is_root(node)) node->_p->_ch[node->_p->_ch[1] == node] = rt;

            if ((node->_ch[ch_idx] = rt->_ch[ch_idx ^ 1])) node->_ch[ch_idx]->_p = node;

            rt->_ch[ch_idx ^ 1] = node;
            rt->_p = std::exchange(node->_p, rt);

            node_type::update(node), node_type::update(rt);
        }

        static void splay(node_pointer_type node) {
            node_type::push(node);
            while (not node_type::is_root(node)) {
                node_pointer_type p = node->_p;
                if (node_type::is_root(p)) {
                    node_type::push(p), node_type::push(node);
                    node_type::rot(p, p->_ch[1] == node);
                } else {
                    node_pointer_type pp = p->_p;
                    node_type::push(pp), node_type::push(p), node_type::push(node);
                    const int idx_pp = pp->_ch[1] == p, idx_p = p->_ch[1] == node;
                    if (idx_p == idx_pp) {
                        node_type::rot(pp, idx_pp), node_type::rot(p, idx_p);
                    } else {
                        node_type::rot(p, idx_p), node_type::rot(pp, idx_pp);
                    }
                }
            }
        }
    };

    template <typename NodeType, typename Derived>
    struct LinkCutTreeBase {
        using derived_tree_type = Derived;

        using node_type = typename NodeType::node_type;
        using node_pointer_type = typename NodeType::node_pointer_type;
        using value_type = typename NodeType::value_type;

        LinkCutTreeBase() = delete;

        static void init_pool(int capacity) {
            _pool = ObjectPool<node_type>(capacity);
        }

        template <typename ...Args>
        static node_pointer_type make_node(Args&&...args) {
            return &(*_pool.alloc() = node_type(std::forward<Args>(args)...));
        }
        static std::vector<node_pointer_type> make_nodes(const std::vector<value_type>& vals) {
            std::vector<node_pointer_type> nodes;
            nodes.reserve(vals.size());
            for (const auto& val : vals) nodes.push_back(make_node(val));
            return nodes;
        }

        static node_pointer_type expose(node_pointer_type node) {
            assert(node);
            node_pointer_type rch = nullptr;
            for (node_pointer_type cur = node; cur; cur = node_type::parent(cur)) {
                node_type::splay(cur);
                node_type::child(cur, 1) = std::exchange(rch, cur);
                node_type::update(cur);
            }
            node_type::splay(node);
            return rch;
        }

        static void link(node_pointer_type ch, node_pointer_type par) {
            derived_tree_type::evert(ch), derived_tree_type::expose(par);
            // check un-connectivity
            if (ch == par or node_type::parent(ch)) assert(false);
            node_type::child(par, 1) = ch;
            node_type::parent(ch) = par;
            node_type::update(par);
        }

        static void cut(node_pointer_type ch) {
            derived_tree_type::expose(ch);
            node_pointer_type par = node_type::child(ch, 0);
            assert(par);
            node_type::parent(par) = node_type::child(ch, 0) = nullptr;
            node_type::update(ch);
        }
        static void cut(node_pointer_type u, node_pointer_type v) {
            derived_tree_type::evert(u);
            derived_tree_type::expose(v);
            // check connectivity
            if (node_type::child(v, 0) != u) assert(false);
            node_type::parent(u) = node_type::child(v, 0) = nullptr;
            node_type::update(v);
        }

        static void evert(node_pointer_type u) {
            derived_tree_type::expose(u);
            node_type::reverse_all(u);
            node_type::push(u);
        }

        static bool is_connected(node_pointer_type u, node_pointer_type v) {
            derived_tree_type::expose(u), derived_tree_type::expose(v);
            return u == v or node_type::parent(u);
        }

        static node_pointer_type lca(node_pointer_type u, node_pointer_type v) {
            derived_tree_type::expose(u);
            node_pointer_type a = derived_tree_type::expose(v);
            return u == v or node_type::parent(u) ? a : nullptr;
        }

        static value_type get(node_pointer_type u) {
            // expose(u);
            return node_type::value(u);
        }
        static void set(node_pointer_type u, const value_type& val) {
            derived_tree_type::update_value(u, [&val](const value_type&) { return val; });
        }
        template <typename Fun>
        static void update_value(node_pointer_type u, Fun&& f) {
            derived_tree_type::expose(u);
            node_type::set_value(u, f(node_type::value(u)));
            node_type::update(u);
        }

        static std::vector<node_pointer_type> path_from_root(node_pointer_type u) {
            std::vector<node_pointer_type> res;
            derived_tree_type::expose(u);
            auto dfs = [&](auto dfs, node_pointer_type cur) -> void {
                node_type::push(cur);
                if (node_type::child(cur, 0)) dfs(dfs, node_type::child(cur, 0));
                res.push_back(cur);
                if (node_type::child(cur, 1)) dfs(dfs, node_type::child(cur, 1));
            };
            dfs(dfs, u);
            return res;
        }
        static std::optional<std::vector<node_pointer_type>> path(node_pointer_type u, node_pointer_type v) {
            derived_tree_type::evert(u);
            derived_tree_type::expose(v);
            if (u == v or node_type::parent(u)) return derived_tree_type::path_from_root(v);
            return std::nullopt;
        }
    
    private:
        static inline ObjectPool<node_type> _pool{};
    };
} // namespace suisen

namespace suisen {
    namespace internal::link_cut_tree {
        template <typename T, T(*op)(T, T), T(*e)(), typename F, T(*mapping)(F, T, int), F(*composition)(F, F), F(*id)(), T(*toggle)(T)>
        struct PathFoldableSplayTreeNodeLazy : SplayTreeNodeBase<T, PathFoldableSplayTreeNodeLazy<T, op, e, F, mapping, composition, id, toggle>> {
            using base_node_type = SplayTreeNodeBase<T, PathFoldableSplayTreeNodeLazy<T, op, e, F, mapping, composition, id, toggle>>;
            
            template <typename, typename>
            friend struct SplayTreeNodeBase;
            template <typename, typename>
            friend struct LinkCutTreeBase;
            template <typename T2, T2(*)(T2, T2), T2(*)(), typename F2, T2(*)(F2, T2, int), F2(*)(F2, F2), F2(*)(), T2(*)(T2)>
            friend struct LinkCutTreePathFoldableLazy;

            using value_type = typename base_node_type::value_type;
            using operator_type = F;
            using node_type = typename base_node_type::node_type;
            using node_pointer_type = typename base_node_type::node_pointer_type;

            explicit PathFoldableSplayTreeNodeLazy(const value_type& val = e()) : base_node_type(val), _sum(val) {}

        protected:
            value_type _sum;
            operator_type _laz = id();

            static value_type sum(node_pointer_type node) {
                return node ? node->_sum : e();
            }
            static void update(node_pointer_type node) {
                base_node_type::update(node);
                node->_sum = op(op(node_type::sum(node_type::child(node, 0)), node_type::value(node)), node_type::sum(node_type::child(node, 1)));
            }
            static void apply_all(node_pointer_type node, const operator_type& f) {
                if (not node) return;
                node->_laz = composition(f, node->_laz);
                node_type::set_value(node, mapping(f, node_type::value(node), 1));
                node->_sum = mapping(f, node->_sum, node_type::size(node));
            }
            static void reverse_all(node_pointer_type node) {
                if (not node) return;
                base_node_type::reverse_all(node);
                node->_sum = toggle(node->_sum);
            }
            static void push(node_pointer_type node) {
                apply_all(node_type::child(node, 0), node->_laz);
                apply_all(node_type::child(node, 1), node->_laz);
                node->_laz = id();
                base_node_type::push(node);
            }
        };

        template <typename T, T(*op)(T, T), T(*e)(), typename F, T(*mapping)(F, T, int), F(*composition)(F, F), F(*id)(), T(*toggle)(T)>
        struct LinkCutTreePathFoldableLazy : LinkCutTreeBase<PathFoldableSplayTreeNodeLazy<T, op, e, F, mapping, composition, id, toggle>, LinkCutTreePathFoldableLazy<T, op, e, F, mapping, composition, id, toggle>> {
            using base_type = LinkCutTreeBase<PathFoldableSplayTreeNodeLazy<T, op, e, F, mapping, composition, id, toggle>, LinkCutTreePathFoldableLazy<T, op, e, F, mapping, composition, id, toggle>>;
            using node_type = typename base_type::node_type;
            using node_pointer_type = typename node_type::node_pointer_type;
            using value_type = typename node_type::value_type;
            using operator_type = typename node_type::operator_type;

            static value_type prod_from_root(node_pointer_type u) {
                base_type::expose(u);
                return node_type::sum(u);
            }
            static value_type prod(node_pointer_type u, node_pointer_type v) {
                base_type::evert(u);
                base_type::expose(v);
                // check connectivity
                if (not (u == v or node_type::parent(u))) assert(false);
                return node_type::sum(v);
            }

            static value_type get(node_pointer_type u) {
                base_type::expose(u);
                return node_type::value(u);
            }

            static void apply_from_root(node_pointer_type u, const operator_type& f) {
                base_type::expose(u);
                node_type::apply_all(u, f);
                node_type::push(u);
            }
            static void apply(node_pointer_type u, node_pointer_type v, const operator_type& f) {
                base_type::evert(u);
                base_type::expose(v);
                assert(u == v or node_type::parent(u)); // check connectivity
                node_type::apply_all(v, f);
                node_type::push(v);
            }
        };
    } // namespace internal::link_cut_tree
    
    using internal::link_cut_tree::LinkCutTreePathFoldableLazy;
} // namespace suisen

long long op(long long x, long long y) {
    return x + y;
}
long long e() {
    return 0;
}
long long mapping(long long f, long long x, int len) {
    return x + f * len;
}
long long composition(long long f, long long g) {
    return f + g;
}
long long id() {
    return 0;
}
long long toggle(long long x) {
    return x;
}

using DynamicTree = suisen::LinkCutTreePathFoldableLazy<long long, op, e, long long, mapping, composition, id, toggle>;

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    int n;
    std::cin >> n;

    DynamicTree::init_pool(n);

    auto nodes = DynamicTree::make_nodes(std::vector<long long>(n));
    for (int i = 0; i < n - 1; ++i) {
        int u, v;
        std::cin >> u >> v;
        --u, --v;
        DynamicTree::link(nodes[u], nodes[v]);
    }

    long long ans = 0;
    int q;
    std::cin >> q;
    while (q --> 0) {
        int u, v;
        std::cin >> u >> v;
        --u, --v;
        DynamicTree::apply(nodes[u], nodes[v], 1);
        ans += DynamicTree::prod(nodes[u], nodes[v]);
    }
    std::cout << ans << '\n';
    return 0;
}