組み合わせ計算の関数の定義. perm(), comb(), fact().
(ugilib/math/combinatorics.hpp)

• View this file on GitHub
• Last update: 2024-04-26 23:52:15+09:00
• Include: #include "ugilib/math/combinatorics.hpp"

Depends on

• ugilib/base/constants.hpp
• ugilib/base/definitions.hpp
• 逆元計算の関数の定義. invmod(). (ugilib/math/invmod.hpp)

Code

#pragma once

/**
 * @file combinatorics.hpp
 * @brief 組み合わせ計算の関数の定義. perm(), comb(), fact().
*/

#include <bits/stdc++.h>
#include "ugilib/base/definitions.hpp"
#include "ugilib/base/constants.hpp"
#include "ugilib/math/invmod.hpp"

using namespace std;

namespace ugilib {
    /**
     * @brief nPrを求める
     * @tparam T 整数型
     * @param n n
     * @param r r
     * @param mod mod
     * @return T nPr
     * @note O(r)
    */
    template <typename T>
    T perm(T n, T r, T mod = ugilib::constants::INF<T>) {
        assert(n >= 0 && r >= 0);
        if (r > n) return 0;
        if (r == 0) return 1;

        T res = 1;
        for (T i = 0; i < r; i++) {
            res = res * (n - i) % mod;
        }
        return res;
    }

    /**
     * @brief n!を求める
     * @tparam T 整数型
     * @param n n
     * @param mod mod
     * @return T n!
     * @note O(n)
     */
    template <typename T>
    T fact(T n, T mod = ugilib::constants::INF<T>) {
        assert(n >= 0);
        return perm(n, n-1, mod);
    }

    /**
     * @brief nCrを求める
     * @tparam T 整数型
     * @param n n
     * @param r r
     * @param mod mod
     * @return T nCr
     * @note O(r)
     */
    template <typename T>
    T comb(T n, T r, T mod = ugilib::constants::INF<T>) {
        assert(n >= 0 && r >= 0);
        if (r > n) return 0;

        r = min(r, n-r);
        if (r == 0) return 1;

        T numer = perm(n, r, mod);
        T denom = fact(r, mod);

        assert(denom != 0);

        return (numer * invmod(denom, mod)) % mod;
    }
}

#line 2 "ugilib/math/combinatorics.hpp"

/**
 * @file combinatorics.hpp
 * @brief 組み合わせ計算の関数の定義. perm(), comb(), fact().
*/

#include <bits/stdc++.h>
#line 2 "ugilib/base/definitions.hpp"

using ll = long long;
using ull = unsigned long long;
using ld = long double;
#define rep(i, n) for(size_t i = 0; i < n; i++)  // rep macro
#define all(v) begin(v), end(v)  // all iterator
#line 3 "ugilib/base/constants.hpp"

namespace ugilib::constants {
    template<typename T>
    inline constexpr T INF = std::numeric_limits<T>::max() / 2;
} // namespace ugilib::constants

const ll INF = ugilib::constants::INF<ll>;
#line 2 "ugilib/math/invmod.hpp"

/**
 * @file invmod.hpp
 * @brief 逆元計算の関数の定義. invmod().
*/

#line 11 "ugilib/math/invmod.hpp"

using namespace std;

namespace ugilib {
    /**
     * @brief 拡張ユークリッドの互除法
     * @tparam T 整数型
     * @param a a
     * @param b b
     * @param x x
     * @param y y
     * @return T gcd(a, b)
     * @note O(log^2(max(a, b)))
     * @note ax + by = gcd(a, b)を満たすx, yを求める
    */
    template <typename T>
    T extended_gcd(T a, T b, T& x, T& y) {
        if (b == 0) {
            x = 1;
            y = 0;
            return a;
        }
        T gcd = extended_gcd(b, a % b, x, y);
        T tmp = x;
        x = y;
        y = tmp - a / b * y;
        return gcd;
    }

    /**
     * @brief aのmodにおける逆元を求める
     * @tparam T 整数型
     * @param a a
     * @param mod mod
     * @return T aのmodにおける逆元
     * @note O(log^2(mod))
     * @note modが素数であり、互いに素であることが前提
    */
    template <typename T>
    T invmod(T a, T mod) {
        T x, y;
        T gcd = extended_gcd(a, mod, x, y);
        assert(gcd == 1);
        return (x % mod + mod) % mod;
    }
}
#line 12 "ugilib/math/combinatorics.hpp"

using namespace std;

namespace ugilib {
    /**
     * @brief nPrを求める
     * @tparam T 整数型
     * @param n n
     * @param r r
     * @param mod mod
     * @return T nPr
     * @note O(r)
    */
    template <typename T>
    T perm(T n, T r, T mod = ugilib::constants::INF<T>) {
        assert(n >= 0 && r >= 0);
        if (r > n) return 0;
        if (r == 0) return 1;

        T res = 1;
        for (T i = 0; i < r; i++) {
            res = res * (n - i) % mod;
        }
        return res;
    }

    /**
     * @brief n!を求める
     * @tparam T 整数型
     * @param n n
     * @param mod mod
     * @return T n!
     * @note O(n)
     */
    template <typename T>
    T fact(T n, T mod = ugilib::constants::INF<T>) {
        assert(n >= 0);
        return perm(n, n-1, mod);
    }

    /**
     * @brief nCrを求める
     * @tparam T 整数型
     * @param n n
     * @param r r
     * @param mod mod
     * @return T nCr
     * @note O(r)
     */
    template <typename T>
    T comb(T n, T r, T mod = ugilib::constants::INF<T>) {
        assert(n >= 0 && r >= 0);
        if (r > n) return 0;

        r = min(r, n-r);
        if (r == 0) return 1;

        T numer = perm(n, r, mod);
        T denom = fact(r, mod);

        assert(denom != 0);

        return (numer * invmod(denom, mod)) % mod;
    }
}

Back to top page