cp-library-cpp

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

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

:heavy_check_mark: library/polynomial/rook_polynomial.hpp

Depends on

Verified with

Code

#ifndef SUISEN_ROOK_POLYNOMIAL
#define SUISEN_ROOK_POLYNOMIAL

#include <algorithm>
#include <cassert>
#include <numeric>
#include <vector>

#include "library/polynomial/convert_to_newton_basis.hpp"
#include "library/datastructure/cartesian_tree.hpp"
#include "library/math/factorial.hpp"

namespace suisen {
    // O(n(log n)^2). returns rook polynomial r s.t. r[k] := # ways to put k non-attacking rooks on a ferrers board
    template <typename FPSType>
    FPSType rook_polynomial_ferrers_board(const std::vector<int> &h) {
        using mint = typename FPSType::value_type;
        assert(std::is_sorted(h.begin(), h.end()));
        const int n = h.size();
        std::vector<FPSType> fs(n);
        for (int i = 0; i < n; ++i) fs[i] = FPSType{ h[i] - i, 1 };
        FPSType f = FPSType::prod(fs);
        std::vector<mint> p(n + 1);
        std::iota(p.begin(), p.end(), 0);
        FPSType r = convert_to_newton_basis(f, p);
        std::reverse(r.begin(), r.end());
        return r;
    }

    // O(n^2 log n). returns rook polynomial r s.t. r[k] := # ways to put k non-attacking rooks on a skyline board
    template <typename FPSType>
    FPSType rook_polynomial_skyline_board(const std::vector<int> &h) {
        using fps = FPSType;
        using mint = typename fps::value_type;

        const int n = h.size();

        factorial<mint> fac(n);

        MinCartesianTreeBuilder ct_builder{h};
        CartesianTree t = ct_builder.build();
        auto dfs = [&](auto dfs, int u, int l, int r) -> fps {
            if (u == t.absent) return { 1 };
            fps f = dfs(dfs, t[u][0], l, u);
            fps g = dfs(dfs, t[u][1], u + 1, r);
            fps fg = f * g; // O(n^2)
            fg.push_back(0);

            const int a = h[u] - (u == t.root ? 0 : h[ct_builder.parent(u)]);
            const int b = r - l;
            assert(int(fg.size()) == b + 1);

            fps s(b + 1), t(b + 1);
            mint binom_a_i = 1;
            for (int i = 0; i <= b; ++i) {
                s[i] = fg[i] * fac.fac(b - i);
                t[i] = binom_a_i;
                binom_a_i *= (a - i) * fac.fac_inv(i + 1) * fac.fac(i);
            }
            fps st = s * t;
            st.resize(b + 1);
            for (int i = 0; i <= b; ++i) {
                st[i] *= fac.fac_inv(b - i);
            }
            return st;
        };

        return dfs(dfs, t.root, 0, n);
    }
} // namespace suisen


#endif // SUISEN_ROOK_POLYNOMIAL
#line 1 "library/polynomial/rook_polynomial.hpp"



#include <algorithm>
#include <cassert>
#include <numeric>
#include <vector>

#line 1 "library/polynomial/convert_to_newton_basis.hpp"



#include <tuple>
#line 6 "library/polynomial/convert_to_newton_basis.hpp"

namespace suisen {
    // Returns b=(b_0,...,b_{N-1}) s.t. f(x) = Sum[i=0,N-1] b_i Prod[j=0,i-1](x - p_j)
    template <typename FPSType>
    std::vector<typename FPSType::value_type> convert_to_newton_basis(const FPSType& f, const std::vector<typename FPSType::value_type>& p) {
        const int n = p.size();
        assert(f.size() == n);

        int m = 1;
        while (m < n) m <<= 1;

        std::vector<FPSType> seg(2 * m);
        for (int i = 0; i < m; ++i) {
            seg[m + i] = { i < n ? -p[i] : 0, 1 };
        }
        for (int i = m - 1; i > 0; --i) {
            if (((i + 1) & -(i + 1)) == (i + 1)) continue; // i = 2^k - 1
            seg[i] = seg[2 * i] * seg[2 * i + 1];
        }

        seg[1] = f;
        for (int i = 1; i < m; ++i) {
            std::tie(seg[2 * i + 1], seg[2 * i]) = seg[i].div_mod(seg[2 * i]);
        }

        std::vector<typename FPSType::value_type> b(n);
        for (int i = 0; i < n; ++i) {
            b[i] = seg[m + i].safe_get(0);
        }
        return b;
    }
} // namespace suisen



#line 1 "library/datastructure/cartesian_tree.hpp"



#include <array>
#line 6 "library/datastructure/cartesian_tree.hpp"
#include <functional>
#line 8 "library/datastructure/cartesian_tree.hpp"

namespace suisen {
    struct CartesianTree : public std::vector<std::array<int, 2>> {
        using base_type = std::vector<std::array<int, 2>>;

        static constexpr int absent = -1;

        const int root;

        CartesianTree() : base_type(), root(0) {}
        CartesianTree(int root, const base_type& g) : base_type(g), root(root) {}
        CartesianTree(int root, base_type&& g) : base_type(std::move(g)), root(root) {}

        auto ranges() const {
            std::vector<std::pair<int, int>> res;
            res.reserve(size());
            auto rec = [&](auto rec, int l, int m, int r) -> void {
                if (m == absent) return;
                const auto& [lm, rm] = (*this)[m];
                rec(rec, l, lm, m), res.emplace_back(l, r), rec(rec, m + 1, rm, r);
            };
            rec(rec, 0, root, size());
            return res;
        }
    };

    template <typename Comparator>
    struct CartesianTreeBuilder {
        CartesianTreeBuilder() = default;
        template <typename RandomAccessibleContainer>
        CartesianTreeBuilder(const RandomAccessibleContainer& a, Comparator comp = Comparator{}) : n(a.size()), comp(comp), par(calc_par(a, comp)) {}

        CartesianTree build() const {
            int root = -1;
            std::vector<std::array<int, 2>> g(n, { CartesianTree::absent, CartesianTree::absent });
            for (int i = 0; i < n; ++i) {
                int p = par[i];
                (p >= 0 ? g[p][p <= i] : root) = i;
            }
            return CartesianTree{ root, std::move(g) };
        }

        template <typename RandomAccessibleContainer>
        static CartesianTree build(const RandomAccessibleContainer& a, Comparator comp = Comparator{}) {
            return CartesianTreeBuilder(a, comp).build();
        }

        int parent(std::size_t i) const {
            assert(i < std::size_t(n));
            return par[i];
        }
        int operator[](std::size_t i) const {
            return parent(i);
        }
    private:
        const int n;
        const Comparator comp;
        const std::vector<int> par;

        template <typename RandomAccessibleContainer>
        static std::vector<int> calc_par(const RandomAccessibleContainer& a, Comparator comp) {
            const int n = a.size();
            std::vector<int> par(n, -1);
            for (int i = 1; i < n; ++i) {
                int p = i - 1, l = i;
                while (p >= 0 and comp(a[i], a[p])) l = p, p = par[p];
                par[l] = i, par[i] = p;
            }
            return par;
        }
    };

    using MinCartesianTreeBuilder = CartesianTreeBuilder<std::less<>>;
    using MaxCartesianTreeBuilder = CartesianTreeBuilder<std::greater<>>;
} // namespace suisen


#line 1 "library/math/factorial.hpp"



#line 6 "library/math/factorial.hpp"

namespace suisen {
    // 引数として与える値に対して、法が十分大きいことを仮定する
    template <typename T, typename U = T>
    struct factorial {
        factorial() = default;
        factorial(int n) { ensure(n); }

        static void ensure(const int n) {
            int sz = _fac.size();
            if (n + 1 <= sz) return;
            int new_size = std::max(n + 1, sz * 2);
            _fac.resize(new_size), _fac_inv.resize(new_size);
            for (int i = sz; i < new_size; ++i) _fac[i] = _fac[i - 1] * i;
            _fac_inv[new_size - 1] = U(1) / _fac[new_size - 1];
            for (int i = new_size - 1; i > sz; --i) _fac_inv[i - 1] = _fac_inv[i] * i;
        }

        T fac(const int i) {
            ensure(i);
            return _fac[i];
        }
        T operator()(int i) {
            return fac(i);
        }
        U fac_inv(const int i) {
            ensure(i);
            return _fac_inv[i];
        }
        // i の逆数
        // i = 0 の場合は assert 違反となる
        U inv(const int i) {
            assert(i > 0);
            ensure(i);
            return _fac_inv[i] * _fac[i - 1];
        }
        U binom(const int n, const int r) {
            if (n < 0 or r < 0 or n < r) return 0;
            ensure(n);
            return _fac[n] * _fac_inv[r] * _fac_inv[n - r];
        }
        // binom(n, r) の逆数
        // binom(n, r) = 0 の場合は assert 違反となる
        U binom_inv(const int n, const int r) {
            assert(r >= 0 and n >= r);
            ensure(n);
            return _fac_inv[n] * _fac[r] * _fac[n - r];
        }
        // n 種類から重複を許して r 個選ぶ場合の数
        // x_1+x_2+...+x_n=r（x_i は非負整数）となる x の個数でもある
        // multichoose(n, r) = binom(n + r - 1, r)
        U multichoose(const int n, const int r) {
            if (n < 0 or r < 0) return 0;
            return r > 0 ? binom(n + r - 1, r) : U(1);
        }
        // n 種類から重複を許して r 個選ぶ場合の数 multichoose(n, r) の逆数
        // x_1+x_2+...+x_n=r（x_i は非負整数）となる x の個数の逆数でもある
        // multichoose(n, r) = binom(n + r - 1, r)
        // multichoose(n, r) = 0 の場合は assert 違反となる
        U multichoose_inv(const int n, const int r) {
            assert(n >= 0 and r >= 0);
            return r > 0 ? binom_inv(n + r - 1, r) : U(1);
        }
        template <typename ...Ds, std::enable_if_t<std::conjunction_v<std::is_integral<Ds>...>, std::nullptr_t> = nullptr>
        U polynom(const int n, const Ds& ...ds) {
            if (n < 0) return 0;
            ensure(n);
            int sumd = 0;
            U res = _fac[n];
            for (int d : { ds... }) {
                if (d < 0 or d > n) return 0;
                sumd += d;
                res *= _fac_inv[d];
            }
            if (sumd > n) return 0;
            res *= _fac_inv[n - sumd];
            return res;
        }
        U perm(const int n, const int r) {
            if (n < 0 or r < 0 or n < r) return 0;
            ensure(n);
            return _fac[n] * _fac_inv[n - r];
        }
        // perm(n, r) の逆数
        // perm(n, r) = 0 の場合は assert 違反となる
        U perm_inv(const int n, const int r) {
            assert(r >= 0 and n >= r);
            ensure(n);
            return _fac_inv[n] * _fac[n - r];
        }
    private:
        static std::vector<T> _fac;
        static std::vector<U> _fac_inv;
    };
    template <typename T, typename U>
    std::vector<T> factorial<T, U>::_fac{ 1 };
    template <typename T, typename U>
    std::vector<U> factorial<T, U>::_fac_inv{ 1 };
} // namespace suisen


#line 12 "library/polynomial/rook_polynomial.hpp"

namespace suisen {
    // O(n(log n)^2). returns rook polynomial r s.t. r[k] := # ways to put k non-attacking rooks on a ferrers board
    template <typename FPSType>
    FPSType rook_polynomial_ferrers_board(const std::vector<int> &h) {
        using mint = typename FPSType::value_type;
        assert(std::is_sorted(h.begin(), h.end()));
        const int n = h.size();
        std::vector<FPSType> fs(n);
        for (int i = 0; i < n; ++i) fs[i] = FPSType{ h[i] - i, 1 };
        FPSType f = FPSType::prod(fs);
        std::vector<mint> p(n + 1);
        std::iota(p.begin(), p.end(), 0);
        FPSType r = convert_to_newton_basis(f, p);
        std::reverse(r.begin(), r.end());
        return r;
    }

    // O(n^2 log n). returns rook polynomial r s.t. r[k] := # ways to put k non-attacking rooks on a skyline board
    template <typename FPSType>
    FPSType rook_polynomial_skyline_board(const std::vector<int> &h) {
        using fps = FPSType;
        using mint = typename fps::value_type;

        const int n = h.size();

        factorial<mint> fac(n);

        MinCartesianTreeBuilder ct_builder{h};
        CartesianTree t = ct_builder.build();
        auto dfs = [&](auto dfs, int u, int l, int r) -> fps {
            if (u == t.absent) return { 1 };
            fps f = dfs(dfs, t[u][0], l, u);
            fps g = dfs(dfs, t[u][1], u + 1, r);
            fps fg = f * g; // O(n^2)
            fg.push_back(0);

            const int a = h[u] - (u == t.root ? 0 : h[ct_builder.parent(u)]);
            const int b = r - l;
            assert(int(fg.size()) == b + 1);

            fps s(b + 1), t(b + 1);
            mint binom_a_i = 1;
            for (int i = 0; i <= b; ++i) {
                s[i] = fg[i] * fac.fac(b - i);
                t[i] = binom_a_i;
                binom_a_i *= (a - i) * fac.fac_inv(i + 1) * fac.fac(i);
            }
            fps st = s * t;
            st.resize(b + 1);
            for (int i = 0; i <= b; ++i) {
                st[i] *= fac.fac_inv(b - i);
            }
            return st;
        };

        return dfs(dfs, t.root, 0, n);
    }
} // namespace suisen
Back to top page