cp-library-cpp

This documentation is automatically generated by online-judge-tools/verification-helper

View the Project on GitHub suisen-cp/cp-library-cpp

:heavy_check_mark: Frequency Table Of Tree Distance
(library/tree/frequency_table_of_tree_distance.hpp)

Frequency Table Of Tree Distance

Depends on

Verified with

Code

#ifndef SUISEN_FREQUENCY_TABLE_OF_TREE_DISTANCE
#define SUISEN_FREQUENCY_TABLE_OF_TREE_DISTANCE

#include <cmath>
#include <iostream>
#include <atcoder/convolution>
#include "library/tree/centroid_decomposition.hpp"

#include "library/util/timer.hpp"

namespace suisen {
    /**
     * @brief Given a tree with N vertices, calculates the number of unordered pair (u, v) s.t. dist(u, v) = k for k = 1, ..., N - 1.
     * @param n size of tree
     * @param edges edges of tree
     * @return std::vector<long long> v : v[k] is the number of unordered pair (u, v) s.t. dist(u, v) = k
     */
    std::vector<long long> frequency_table_of_tree_distance(const size_t n, const std::vector<std::pair<int, int>>& edges) {
        static constexpr long long MOD1 = 998244353;
        static constexpr long long MOD2 = 754974721;
        static constexpr long long INV_MOD1 = atcoder::internal::inv_gcd(MOD1, MOD2).second;

        using mint1 = atcoder::static_modint<MOD1>;
        using mint2 = atcoder::static_modint<MOD2>;

        auto garner = [](mint1 x1, mint2 x2) -> long long {
            long long v1 = x1.val(), v2 = x2.val();
            return mint2((v2 - v1) * INV_MOD1).val() * MOD1 + v1;
        };

        std::vector<long long> res(n);

        CentroidDecompositionUnweighted g(UnweightedGraph::create_undirected_graph(n, edges));

        const int garner_threshold = ::sqrt(MOD1);

        for (size_t root = 0; root < n; ++root) {
            std::vector<std::vector<std::pair<int, int>>> d = g.collect(root, 0);
            for (size_t i = 0; i < d.size(); ++i) {
                int max_dep = 0;
                for (auto [v, w] : d[i]) max_dep = std::max(max_dep, w);
                std::vector<int> f(max_dep + 1);
                for (auto [v, w] : d[i]) ++f[w];

                std::vector<int> sq_f1 = atcoder::convolution<mint1::mod()>(f, f);
                int coef = i == 0 ? +1 : -1;
                if (d[i].size() <= garner_threshold) {
                    for (size_t i = 0; i < std::min(n, sq_f1.size()); ++i) {
                        res[i] += coef * sq_f1[i];
                    }
                } else {
                    std::vector<int> sq_f2 = atcoder::convolution<mint2::mod()>(f, f);
                    for (size_t i = 0; i < std::min(n, sq_f1.size()); ++i) {
                        res[i] += coef * garner(sq_f1[i], sq_f2[i]);
                    }
                }
            }
        }

        for (size_t i = 1; i < n; ++i) {
            res[i] >>= 1;
        }
        return res;
    }
} // namespace suisen


#endif // SUISEN_FREQUENCY_TABLE_OF_TREE_DISTANCE
#line 1 "library/tree/frequency_table_of_tree_distance.hpp"



#include <cmath>
#include <iostream>
#include <atcoder/convolution>
#line 1 "library/tree/centroid_decomposition.hpp"



#include <deque>
#include <limits>
#include <queue>
#include <tuple>
#include <vector>
#line 1 "library/graph/csr_graph.hpp"



#include <algorithm>
#include <cassert>
#line 7 "library/graph/csr_graph.hpp"
#include <optional>
#include <type_traits>
#line 10 "library/graph/csr_graph.hpp"
#include <utility>
#line 12 "library/graph/csr_graph.hpp"

namespace suisen {
    namespace internal::csr_graph { struct graph_base_tag {}; }
    struct directed_graph_tag : internal::csr_graph::graph_base_tag {};
    struct undirected_graph_tag : internal::csr_graph::graph_base_tag {};
    template <typename T>
    struct is_graph_tag { static constexpr bool value = std::is_base_of_v<internal::csr_graph::graph_base_tag, T>; };
    template <typename T>
    constexpr bool is_graph_tag_v = is_graph_tag<T>::value;

    template <typename WeightType = void>
    struct Graph {
        template <typename GraphTag, typename, std::enable_if_t<is_graph_tag_v<GraphTag>, std::nullptr_t>>
        friend struct GraphBuilder;

        using weight_type = WeightType;
        static constexpr bool weighted = std::negation_v<std::is_same<weight_type, void>>;

        using weight_type_or_1 = std::conditional_t<weighted, weight_type, int>;

        using input_edge_type = std::conditional_t<weighted, std::tuple<int, int, weight_type>, std::pair<int, int>>;
    private:
        using internal_edge_type = std::conditional_t<weighted, std::pair<int, weight_type>, int>;
        struct Edge : public internal_edge_type {
            using internal_edge_type::internal_edge_type;
            operator int() const { return std::get<0>(*this); }
        };
    public:
        using edge_type = std::conditional_t<weighted, Edge, int>;
    private:
        struct AdjacentList {
            friend struct Graph;

            using value_type = edge_type;
            using iterator = typename std::vector<value_type>::iterator;
            using const_iterator = typename std::vector<value_type>::const_iterator;
            using reverse_iterator = typename std::vector<value_type>::reverse_iterator;
            using const_reverse_iterator = typename std::vector<value_type>::const_reverse_iterator;

            AdjacentList() = default;

            int size() const { return _siz; }
            bool empty() const { return _siz == 0; }
            int capacity() const { return _cap; }

            value_type& operator[](int i) { return *(begin() + i); }
            const value_type& operator[](int i) const { return *(cbegin() + i); }
            value_type& at(uint32_t i) { assert(i < _siz); return *(begin() + i); }
            const value_type& at(uint32_t i) const { assert(i < _siz); return *(cbegin() + i); }

            value_type* data() { return _g->_edges.data() + _offset; }
            const value_type* data() const { return _g->_edges.data() + _offset; }

            iterator begin() const { return _g->_edges.begin() + _offset; }
            iterator end() const { return begin() + _siz; }
            const_iterator cbegin() const { return _g->_edges.cbegin() + _offset; }
            const_iterator cend() const { return cbegin() + _siz; }
            reverse_iterator rbegin() const { return _g->_edges.rbegin() + (_g->_edges.size() - (_offset + _siz)); }
            reverse_iterator rend() const { return rbegin() + _siz; }
            const_reverse_iterator crbegin() const { return _g->_edges.crbegin() + (_g->_edges.size() - (_offset + _siz)); }
            const_reverse_iterator crend() const { return crbegin() + _siz; }

            void erase(const_iterator pos) {
                erase(pos, std::next(pos));
            }
            void erase(const_iterator first, const_iterator last) {
                const int num = last - first, k = first - cbegin();
                assert(num >= 0);
                if (num == 0) return;
                assert(0 <= k and k <= _siz - num);
                std::move(begin() + k + num, end(), begin() + k);
                _siz -= num;
            }
            void pop_back() {
                assert(_siz);
                --_siz;
            }
            void clear() { _siz = 0; }

            const value_type& back() const { return *--cend(); }
            value_type& back() { return *--end(); }
            const value_type& front() const { return *cbegin(); }
            value_type& front() { return *begin(); }

            void push_back(const value_type& x) {
                ++_siz;
                assert(_siz <= _cap);
                back() = x;
            }
            template <typename ...Args>
            void emplace_back(Args &&...args) {
                ++_siz;
                assert(_siz <= _cap);
                back() = value_type(std::forward<Args>(args)...);
            }

            void insert(const_iterator pos, const value_type& x) {
                emplace(pos, x);
            }
            void insert(const_iterator pos, int num, const value_type& x) {
                const int k = pos - cbegin();
                assert(0 <= k and k <= _siz);
                std::fill(begin() + k, shift_back(begin() + k, num), x);
            }
            template <class RandomAccessIterator>
            auto insert(const_iterator pos, RandomAccessIterator first, RandomAccessIterator last) -> decltype(*first++, last - first, void()) {
                const int num = last - first, k = pos - cbegin();
                assert(0 <= k and k <= _siz);
                shift_back(begin() + k, num);
                std::copy(first, last, begin() + k);
            }
            void insert(const_iterator pos, std::initializer_list<value_type> il) { insert(pos, il.begin(), il.end()); }
            template <typename ...Args>
            void emplace(const_iterator pos, Args &&...args) {
                const int k = pos - cbegin();
                assert(0 <= k and k <= _siz);
                *--shift_back(begin() + k) = value_type(std::forward<Args>(args)...);
            }
        private:
            mutable Graph* _g;
            int _cap;
            int _offset;
            int _siz;

            iterator shift_back(iterator pos, int num = 1) {
                _siz += num;
                assert(_siz <= _cap);
                return std::move_backward(pos, end() - num, end());
            }
        };
    public:
        using adjacent_list = AdjacentList;

        Graph() = default;

        template <typename GraphTag, std::enable_if_t<is_graph_tag_v<GraphTag>, std::nullptr_t> = nullptr>
        Graph(const int n, const std::vector<input_edge_type>& edges, GraphTag, std::vector<int> cap = {}) : _n(n), _adj(_n) {
            static constexpr bool undirected = std::is_same_v<undirected_graph_tag, GraphTag>;

            for (const auto& e : edges) {
                const int u = std::get<0>(e);
                ++_adj[u]._siz;
                if constexpr (undirected) {
                    const int v = std::get<1>(e);
                    ++_adj[v]._siz;
                }
            }
            if (cap.empty()) cap.resize(_n, std::numeric_limits<int>::max());
            int edge_num = 0;
            for (int i = 0; i < _n; ++i) {
                _adj[i]._g = this;
                _adj[i]._cap = std::min(_adj[i]._siz, cap[i]);
                _adj[i]._offset = edge_num;
                edge_num += _adj[i]._siz;
            }
            _edges.resize(edge_num);
            std::vector<typename std::vector<edge_type>::iterator> ptr(_n);
            for (int i = 0; i < _n; ++i) ptr[i] = _adj[i].begin();
            for (const auto& e : edges) {
                const int u = std::get<0>(e);
                const int v = std::get<1>(e);
                if constexpr (weighted) {
                    const weight_type& w = std::get<2>(e);
                    *ptr[u]++ = { v, w };
                    if constexpr (undirected) *ptr[v]++ = { u, w };
                } else {
                    *ptr[u]++ = v;
                    if constexpr (undirected) *ptr[v]++ = u;
                }
            }
        }
        Graph(const std::vector<std::vector<edge_type>>& g) : Graph(g.size(), make_edges(g), directed_graph_tag{}) {}

        static Graph create_directed_graph(const int n, const std::vector<input_edge_type>& edges, const std::vector<int>& cap = {}) {
            return Graph(n, edges, directed_graph_tag{}, cap);
        }
        static Graph create_undirected_graph(const int n, const std::vector<input_edge_type>& edges, const std::vector<int>& cap = {}) {
            return Graph(n, edges, undirected_graph_tag{}, cap);
        }

        adjacent_list& operator[](int i) {
            _adj[i]._g = this;
            return _adj[i];
        }
        const adjacent_list& operator[](int i) const {
            _adj[i]._g = const_cast<Graph*>(this);
            return _adj[i];
        }

        int size() const {
            return _n;
        }

        void shrink_to_fit() {
            int edge_num = 0;
            for (const auto& l : _adj) edge_num += l.size();

            std::vector<edge_type> new_edges(edge_num);
            auto it = new_edges.begin();
            for (int i = 0; i < _n; ++i) {
                int nl = it - new_edges.begin();
                it = std::move(_adj[i].begin(), _adj[i].end(), it);
                _adj[i]._offset = nl;
                _adj[i]._cap = _adj[i]._siz;
            }
            _edges.swap(new_edges);
        }

        static weight_type_or_1 get_weight(const edge_type& edge) {
            if constexpr (weighted) return std::get<1>(edge);
            else return 1;
        }

        Graph reversed(const std::vector<int>& cap = {}) const {
            std::vector<input_edge_type> edges;
            for (int i = 0; i < _n; ++i) {
                for (const auto& edge : (*this)[i]) {
                    if constexpr (weighted) edges.emplace_back(std::get<0>(edge), i, std::get<1>(edge));
                    else edges.emplace_back(edge, i);
                }
            }
            return Graph(_n, std::move(edges), directed_graph_tag{}, cap);
        }

        struct DFSTree {
            std::vector<int> par;
            std::vector<int> pre_ord, pst_ord;
            Graph tree, back;
        };

        DFSTree dfs_tree(int start = 0) const {
            std::vector<input_edge_type> tree_edge, back_edge;

            std::vector<int> pre(_n), pst(_n);
            auto pre_it = pre.begin(), pst_it = pst.begin();

            std::vector<int> eid(_n, -1), par(_n, -2);
            std::vector<std::optional<weight_type_or_1>> par_w(_n, std::nullopt);
            for (int i = 0; i < _n; ++i) {
                int cur = (start + i) % _n;
                if (par[cur] != -2) continue;
                par[cur] = -1;
                while (cur >= 0) {
                    ++eid[cur];
                    if (eid[cur] == 0) *pre_it++ = cur;
                    if (eid[cur] == _adj[cur].size()) {
                        *pst_it++ = cur;
                        cur = par[cur];
                    } else {
                        const auto &e = _adj[cur][eid[cur]];
                        weight_type_or_1 w = get_weight(e);
                        int nxt = e;
                        if (par[nxt] == -2) {
                            tree_edge.emplace_back(make_edge(cur, e));
                            par[nxt] = cur;
                            par_w[nxt] = std::move(w);
                            cur = nxt;
                        } else if (eid[nxt] != _adj[nxt].size()) {
                            if (par[cur] != nxt or par_w[cur] != w or not std::exchange(par_w[cur], std::nullopt).has_value()) {
                                back_edge.emplace_back(make_edge(cur, e));
                            }
                        }
                    }
                }
            }
            Graph tree = create_directed_graph(_n, tree_edge);
            Graph back = create_directed_graph(_n, back_edge);
            return DFSTree{ std::move(par), std::move(pre), std::move(pst), std::move(tree), std::move(back) };
        }

    private:
        int _n;
        std::vector<adjacent_list> _adj;
        std::vector<edge_type> _edges;

        static std::vector<input_edge_type> make_edges(const std::vector<std::vector<edge_type>>& g) {
            const int n = g.size();
            std::vector<input_edge_type> edges;
            for (int i = 0; i < n; ++i) for (const auto& e : g[i]) {
                edges.emplace_back(make_edge(i, e));
            }
            return edges;
        }
        static input_edge_type make_edge(int i, const edge_type& e) {
            if constexpr (weighted) return { i, std::get<0>(e), std::get<1>(e) };
            else return { i, e };
        }
    };

    template <typename GraphTag>
    Graph(int, std::vector<std::pair<int, int>>, GraphTag, std::vector<int> = {})->Graph<void>;
    template <typename WeightType, typename GraphTag>
    Graph(int, std::vector<std::tuple<int, int, WeightType>>, GraphTag, std::vector<int> = {})->Graph<WeightType>;

    Graph(std::vector<std::vector<int>>)->Graph<void>;
    template <typename WeightType>
    Graph(std::vector<std::vector<std::pair<int, WeightType>>>)->Graph<WeightType>;

    template <typename GraphTag, typename WeightType = void,
        std::enable_if_t<is_graph_tag_v<GraphTag>, std::nullptr_t> = nullptr>
    struct GraphBuilder {
        using graph_tag = GraphTag;
        using weight_type = WeightType;
        using edge_type = typename Graph<weight_type>::input_edge_type;

        GraphBuilder(int n = 0) : _n(n) {}

        void add_edge(const edge_type& edge) {
            check_not_moved();
            _edges.push_back(edge);
        }
        template <typename ...Args>
        void emplace_edge(Args &&...args) {
            check_not_moved();
            _edges.emplace_back(std::forward<Args>(args)...);
        }
        template <typename EdgeContainer, std::enable_if_t<std::is_constructible_v<edge_type, typename EdgeContainer::value_type>, std::nullptr_t> = nullptr>
        void add_edges(const EdgeContainer& edges) {
            for (const auto& edge : edges) add_edge(edge);
        }

        template <bool move_edges = true>
        Graph<weight_type> build() {
            if constexpr (move_edges) {
                _moved = true;
                return Graph<weight_type>(_n, std::move(_edges), graph_tag{});
            } else {
                return Graph<weight_type>(_n, _edges, graph_tag{});
            }
        }
        Graph<weight_type> build_without_move() {
            return build<false>();
        }

        static Graph<weight_type> build(const int n, const std::vector<edge_type>& edges) {
            GraphBuilder builder(n);
            builder.add_edges(edges);
            return builder.build();
        }
    private:
        int _n;
        std::vector<edge_type> _edges;
        bool _moved = false;

        void check_not_moved() {
            if (not _moved) return;
            std::cerr << "[\033[31mERROR\033[m] Edges are already moved. If you want to add edges after calling build() and build another graph, you should use build_without_move() instead." << std::endl;
            assert(false);
        }
    };
    template <typename WeightType = void>
    using DirectedGraphBuilder = GraphBuilder<directed_graph_tag, WeightType>;
    template <typename WeightType = void>
    using UndirectedGraphBuilder = GraphBuilder<undirected_graph_tag, WeightType>;

    template <typename Weight, std::enable_if_t<std::negation_v<std::is_same<Weight, void>>, std::nullptr_t> = nullptr>
    using WeightedGraph = Graph<Weight>;
    using UnweightedGraph = Graph<void>;

    template <typename T>
    struct is_weighted_graph { static constexpr bool value = false; };
    template <typename WeightType>
    struct is_weighted_graph<Graph<WeightType>> { static constexpr bool value = Graph<WeightType>::weighted; };
    template <typename T>
    constexpr bool is_weighted_graph_v = is_weighted_graph<T>::value;

    template <typename T>
    struct is_unweighted_graph { static constexpr bool value = false; };
    template <typename WeightType>
    struct is_unweighted_graph<Graph<WeightType>> { static constexpr bool value = not Graph<WeightType>::weighted; };
    template <typename T>
    constexpr bool is_unweighted_graph_v = is_unweighted_graph<T>::value;
} // namespace suisen


#line 10 "library/tree/centroid_decomposition.hpp"

namespace suisen {
    namespace internal {
        template <typename WeightType = void>
        struct CentroidDecomposition : Graph<WeightType> {
            friend struct CentroidDecompositionUnweighted;
            template <typename WeightType_, std::enable_if_t<not std::is_same_v<WeightType_, void>, std::nullptr_t>>
            friend struct CentroidDecompositionWeighted;

            using graph_type = Graph<WeightType>;
            using weight_type = WeightType;

            CentroidDecomposition(const graph_type& g) : graph_type(g), n(this->size()), cpar(n, -1), cdep(n, std::numeric_limits<int>::max()), csiz(n) {
                build();
            }

            int dct_parent(int i) const { return cpar[i]; }
            int dct_depth(int i) const { return cdep[i]; }
            int dct_size(int i) const { return csiz[i]; }

        private:
            int n;
            std::vector<int> cpar;
            std::vector<int> cdep;
            std::vector<int> csiz;

            void build() {
                std::vector<int> eid(n, 0);

                cpar[0] = -1, csiz[0] = n;
                std::deque<std::tuple<int, int>> dq{ { 0, 0 } };

                while (dq.size()) {
                    const auto [r, dep] = dq.front();
                    const int siz = csiz[r], prev_ctr = cpar[r];
                    dq.pop_front();

                    int c = -1;
                    eid[r] = 0, csiz[r] = 1, cpar[r] = -1;
                    for (int cur = r;;) {
                        for (const int edge_num = int((*this)[cur].size());;) {
                            if (eid[cur] == edge_num) {
                                if (csiz[cur] * 2 > siz) {
                                    c = cur;
                                } else {
                                    const int nxt = cpar[cur];
                                    csiz[nxt] += csiz[cur];
                                    cur = nxt;
                                }
                                break;
                            }
                            const int nxt = (*this)[cur][eid[cur]++];
                            if (cdep[nxt] >= dep and nxt != cpar[cur]) {
                                eid[nxt] = 0, csiz[nxt] = 1, cpar[nxt] = cur;
                                cur = nxt;
                                break;
                            }
                        }
                        if (c >= 0) break;
                    }
                    for (int v : (*this)[c]) if (cdep[v] >= dep) {
                        if (cpar[c] == v) cpar[v] = c, csiz[v] = siz - csiz[c];
                        dq.emplace_back(v, dep + 1);
                    }
                    cpar[c] = prev_ctr, cdep[c] = dep, csiz[c] = siz;
                }
            }
        };

        struct CentroidDecompositionUnweighted : internal::CentroidDecomposition<void> {
            using base_type = internal::CentroidDecomposition<void>;
            using base_type::base_type;

            std::vector<std::vector<std::pair<int, int>>> collect(int root, int root_val = 0) const {
                std::vector<std::vector<std::pair<int, int>>> res{ { { root, root_val } } };
                for (int sub_root : (*this)[root]) if (this->cdep[sub_root] > this->cdep[root]) {
                    res.emplace_back();
                    std::deque<std::tuple<int, int, int>> dq{ { sub_root, root, root_val + 1 } };
                    while (dq.size()) {
                        auto [u, p, w] = dq.front();
                        dq.pop_front();
                        res.back().emplace_back(u, w);
                        for (int v : (*this)[u]) if (v != p and this->cdep[v] > this->cdep[root]) {
                            dq.emplace_back(v, u, w + 1);
                        }
                    }
                    std::copy(res.back().begin(), res.back().end(), std::back_inserter(res.front()));
                }
                return res;
            }
        };

        template <typename WeightType, std::enable_if_t<not std::is_same_v<WeightType, void>, std::nullptr_t> = nullptr>
        struct CentroidDecompositionWeighted : internal::CentroidDecomposition<WeightType> {
            using base_type = internal::CentroidDecomposition<WeightType>;
            using base_type::base_type;

            using weight_type = typename base_type::weight_type;

            template <typename Op, std::enable_if_t<std::is_invocable_r_v<weight_type, Op, weight_type, weight_type>, std::nullptr_t> = nullptr>
            std::vector<std::vector<std::pair<int, weight_type>>> collect(int root, Op op, weight_type root_val) const {
                std::vector<std::vector<std::pair<int, weight_type>>> res{ { { root, root_val } } };
                for (auto [sub_root, ew] : (*this)[root]) if (this->cdep[sub_root] > this->cdep[root]) {
                    res.emplace_back();
                    std::deque<std::tuple<int, int, weight_type>> dq{ { sub_root, root, op(root_val, ew) } };
                    while (dq.size()) {
                        auto [u, p, w] = dq.front();
                        dq.pop_front();
                        res.back().emplace_back(u, w);
                        for (auto [v, ew] : (*this)[u]) if (v != p and this->cdep[v] > this->cdep[root]) {
                            dq.emplace_back(v, u, op(w, ew));
                        }
                    }
                    std::copy(res.back().begin(), res.back().end(), std::back_inserter(res.front()));
                }
                return res;
            }
        };
    }

    using CentroidDecompositionUnweighted = internal::CentroidDecompositionUnweighted;
    template <typename WeightType, std::enable_if_t<not std::is_same_v<WeightType, void>, std::nullptr_t> = nullptr>
    using CentroidDecompositionWeighted = internal::CentroidDecompositionWeighted<WeightType>;

} // namespace suisen


#line 8 "library/tree/frequency_table_of_tree_distance.hpp"

#line 1 "library/util/timer.hpp"



#include <chrono>

namespace suisen {
    struct Timer {
        using minutes_t = std::chrono::minutes;
        using seconds_t = std::chrono::seconds;
        using milliseconds_t = std::chrono::milliseconds;
        using microseconds_t = std::chrono::microseconds;
        using nanoseconds_t = std::chrono::nanoseconds;

        Timer() { start(); }

        void start() {
            _start = std::chrono::system_clock::now();
        }
        template <typename TimeUnit = std::chrono::milliseconds>
        double elapsed() const {
            return std::chrono::duration_cast<TimeUnit>(std::chrono::system_clock::now() - _start).count();
        }

        template <typename TimeUnit = std::chrono::milliseconds, typename Proc>
        static double measure(Proc &&proc) {
            Timer timer;
            proc();
            return timer.elapsed<TimeUnit>();
        }
    private:
        decltype(std::chrono::system_clock::now()) _start;
    };
} // namespace suisen



#line 10 "library/tree/frequency_table_of_tree_distance.hpp"

namespace suisen {
    /**
     * @brief Given a tree with N vertices, calculates the number of unordered pair (u, v) s.t. dist(u, v) = k for k = 1, ..., N - 1.
     * @param n size of tree
     * @param edges edges of tree
     * @return std::vector<long long> v : v[k] is the number of unordered pair (u, v) s.t. dist(u, v) = k
     */
    std::vector<long long> frequency_table_of_tree_distance(const size_t n, const std::vector<std::pair<int, int>>& edges) {
        static constexpr long long MOD1 = 998244353;
        static constexpr long long MOD2 = 754974721;
        static constexpr long long INV_MOD1 = atcoder::internal::inv_gcd(MOD1, MOD2).second;

        using mint1 = atcoder::static_modint<MOD1>;
        using mint2 = atcoder::static_modint<MOD2>;

        auto garner = [](mint1 x1, mint2 x2) -> long long {
            long long v1 = x1.val(), v2 = x2.val();
            return mint2((v2 - v1) * INV_MOD1).val() * MOD1 + v1;
        };

        std::vector<long long> res(n);

        CentroidDecompositionUnweighted g(UnweightedGraph::create_undirected_graph(n, edges));

        const int garner_threshold = ::sqrt(MOD1);

        for (size_t root = 0; root < n; ++root) {
            std::vector<std::vector<std::pair<int, int>>> d = g.collect(root, 0);
            for (size_t i = 0; i < d.size(); ++i) {
                int max_dep = 0;
                for (auto [v, w] : d[i]) max_dep = std::max(max_dep, w);
                std::vector<int> f(max_dep + 1);
                for (auto [v, w] : d[i]) ++f[w];

                std::vector<int> sq_f1 = atcoder::convolution<mint1::mod()>(f, f);
                int coef = i == 0 ? +1 : -1;
                if (d[i].size() <= garner_threshold) {
                    for (size_t i = 0; i < std::min(n, sq_f1.size()); ++i) {
                        res[i] += coef * sq_f1[i];
                    }
                } else {
                    std::vector<int> sq_f2 = atcoder::convolution<mint2::mod()>(f, f);
                    for (size_t i = 0; i < std::min(n, sq_f1.size()); ++i) {
                        res[i] += coef * garner(sq_f1[i], sq_f2[i]);
                    }
                }
            }
        }

        for (size_t i = 1; i < n; ++i) {
            res[i] >>= 1;
        }
        return res;
    }
} // namespace suisen
Back to top page