Combination (Large)
(math/combination_large.hpp)

• View this file on GitHub
• Last update: 2024-01-08 02:22:28+09:00
• Include: #include "math/combination_large.hpp"

Depends on

• Disjoint Sparse Table (data-structure/disjoint_sparse_table.hpp)

Code

#pragma once
#include <vector>

#include "../data-structure/disjoint_sparse_table.hpp"

/**
 * @brief Combination (Large)
 */
template <typename mint>
class CombinationLarge {
   public:
    CombinationLarge() = default;
    CombinationLarge(int N, long long M)
        : N(N), M(M), fact_(N + 1), fact_inv_(N + 1) {
        fact_[0] = 1;
        for (int i = 1; i <= N; ++i) fact_[i] = fact_[i - 1] * i;
        fact_inv_[N] = fact_[N].inv();
        for (int i = N; i > 0; --i) fact_inv_[i - 1] = fact_inv_[i] * i;

        std::vector<mint> num(N + 1);
        for (int i = 0; i <= N; ++i) {
            num[i] = M - i;
        }
        st = DisjointSparseTable<MulMonoid>(num);
    }

    mint perm(long long n, int k) const {
        assert(M - N + k - 1 <= n && n <= M);
        if (k < 0 || n < k) return 0;
        return st.fold(M - n, M - n + k);
    }

    mint comb(long long n, int k) const {
        assert(M - N + k - 1 <= n && n <= M);
        if (k < 0 || n < k) return 0;
        return st.fold(M - n, M - n + k) * fact_inv_[k];
    }

    mint fact(int n) const { return fact_[n]; }
    mint fact_inv(int n) const { return fact_inv_[n]; }

   private:
    struct MulMonoid {
        using T = mint;
        static T id() { return 1; }
        static T op(T a, T b) { return a * b; }
    };

    int N;
    long long M;
    std::vector<mint> fact_, fact_inv_;
    DisjointSparseTable<MulMonoid> st;
};

#line 2 "math/combination_large.hpp"
#include <vector>

#line 2 "data-structure/disjoint_sparse_table.hpp"
#include <algorithm>
#include <bit>
#line 5 "data-structure/disjoint_sparse_table.hpp"

template <typename M>
class DisjointSparseTable {
    using T = M::T;

   public:
    DisjointSparseTable() = default;
    explicit DisjointSparseTable(const std::vector<T>& v) {
        const int n = v.size();
        const int b = std::bit_width((unsigned int)n);
        lookup.resize(b + 1, std::vector<T>(n));
        std::ranges::copy(v, lookup[0].begin());
        for (int i = 1; i <= b; ++i) {
            int len = 1 << i;
            for (int l = 0; l + len / 2 < n; l += len) {
                int m = l + len / 2;
                lookup[i][m - 1] = v[m - 1];
                for (int j = m - 2; j >= l; --j) {
                    lookup[i][j] = M::op(lookup[i][j + 1], v[j]);
                }
                lookup[i][m] = v[m];
                for (int j = m + 1; j < std::min(l + len, n); ++j) {
                    lookup[i][j] = M::op(lookup[i][j - 1], v[j]);
                }
            }
        }
    }

    T fold(int l, int r) const {
        if (l == r) return M::id();
        if (r - l == 1) return lookup[0][l];
        int i = std::bit_width((unsigned int)(l ^ (r - 1)));
        return M::op(lookup[i][l], lookup[i][r - 1]);
    }

   private:
    std::vector<std::vector<T>> lookup;
};
#line 5 "math/combination_large.hpp"

/**
 * @brief Combination (Large)
 */
template <typename mint>
class CombinationLarge {
   public:
    CombinationLarge() = default;
    CombinationLarge(int N, long long M)
        : N(N), M(M), fact_(N + 1), fact_inv_(N + 1) {
        fact_[0] = 1;
        for (int i = 1; i <= N; ++i) fact_[i] = fact_[i - 1] * i;
        fact_inv_[N] = fact_[N].inv();
        for (int i = N; i > 0; --i) fact_inv_[i - 1] = fact_inv_[i] * i;

        std::vector<mint> num(N + 1);
        for (int i = 0; i <= N; ++i) {
            num[i] = M - i;
        }
        st = DisjointSparseTable<MulMonoid>(num);
    }

    mint perm(long long n, int k) const {
        assert(M - N + k - 1 <= n && n <= M);
        if (k < 0 || n < k) return 0;
        return st.fold(M - n, M - n + k);
    }

    mint comb(long long n, int k) const {
        assert(M - N + k - 1 <= n && n <= M);
        if (k < 0 || n < k) return 0;
        return st.fold(M - n, M - n + k) * fact_inv_[k];
    }

    mint fact(int n) const { return fact_[n]; }
    mint fact_inv(int n) const { return fact_inv_[n]; }

   private:
    struct MulMonoid {
        using T = mint;
        static T id() { return 1; }
        static T op(T a, T b) { return a * b; }
    };

    int N;
    long long M;
    std::vector<mint> fact_, fact_inv_;
    DisjointSparseTable<MulMonoid> st;
};

Back to top page