algo/math/combo.h

• View this file on GitHub
• Last update: 2026-09-28 20:05:44-04:00

Depends on

• algo/common.h
• algo/math/common.h

Required by

• main.cpp

Verified with

• verify/math/combo.test.cpp

Code

#pragma once
#include "algo/common.h"
#include "algo/math/common.h"

namespace algo::math {

// Factorials and binomials mod p. Tables grow as needed.
//   combo<mint> c;
//   c.cmb(n, r);   // n choose r
//   c.perm(n, r);  // n! / (n - r)!
template <typename T>
struct combo {
    explicit combo(int n = 0) {
        if (n > 0) fact(n), inv_fact(n);
    }

    T fact(int n) {
        if (n >= (int)f.size()) {
            assert(n < mod());
            if (f.empty()) f.push_back(T(1));
            int m = grow_to(n, (int)f.size());
            f.reserve(m);
            for (int i = (int)f.size(); i < m; i++) {
                f.push_back(f.back() * T(i));
            }
        }
        return f[n];
    }
    T inv_fact(int n) {
        if (n >= (int)inv_f.size()) {
            assert(n < mod());
            if (inv_f.empty()) inv_f.push_back(T(1));
            int lo = (int)inv_f.size(), m = grow_to(n, lo);
            inv_f.resize(m);
            inv_f[m - 1] = T(1) / fact(m - 1);
            for (int i = m - 2; i >= lo; i--) {
                inv_f[i] = inv_f[i + 1] * T(i + 1);
            }
        }
        return inv_f[n];
    }
    T cmb(int n, int r) {
        if (r < 0 || r > n) {
            return T(0);
        } else {
            return fact(n) * inv_fact(r) * inv_fact(n - r);
        }
    }
    T perm(int n, int r) {
        if (r < 0 || r > n) {
            return T(0);
        } else {
            return fact(n) * inv_fact(n - r);
        }
    }

private:
    std::vector<T> f, inv_f;

    static int mod() {
        if constexpr (requires { T::mod(); }) {
            return T::mod();
        } else {
            return std::numeric_limits<int>::max();
        }
    }
    // New table size. Capped at mod, since n! = 0 for n >= mod.
    static int grow_to(int n, int cur) {
        return std::min<int64_t>(std::max<int64_t>(n + 1, 2LL * cur), mod());
    }
};

} // namespace algo::math

#line 2 "algo/common.h"
#ifndef PREPROCESS
#include <bits/stdc++.h>
#include <cassert>
#endif

// Declared here so `using namespace algo;` works with no other includes.
namespace algo {}
#line 3 "algo/math/common.h"

namespace algo::math {

// x mod m, in [0, m) even for negative x
constexpr int64_t safe_mod(int64_t x, int64_t m) {
    x %= m;
    if (x < 0) x += m;
    return x;
}

// Fast x % m for an m only known at run time.
//   barrett b(m);
//   b.reduce(x);  // x % m
struct barrett {
    constexpr explicit barrett(uint64_t _m) : m(_m), im(-1ULL / _m) {
        assert(1 <= _m);
    }
    uint64_t mod() const {
        return m;
    };
    uint64_t reduce(uint64_t a) const {
        uint64_t q = (uint64_t)((__uint128_t(im) * a) >> 64);
        uint64_t r = a - q * m;
        return r - (r >= m) * m;
    }

private:
    uint64_t m, im;
};

// Rounded-up and rounded-down a / b. Correct for negatives.
constexpr int64_t c_div(int64_t a, int64_t b) {
    return a / b + ((a ^ b) > 0 && a % b);
}
constexpr int64_t f_div(int64_t a, int64_t b) {
    return a / b - ((a ^ b) < 0 && a % b);
}

// x^n by repeated squaring.
//   bpow(mint(2), n);
//   bpow(m, n, identity, mat_mul);  // any op, with its identity
constexpr auto bpow(auto const &x, auto n, auto const &one, auto op) {
    if (n == 0) {
        return one;
    } else {
        auto t = bpow(x, n / 2, one, op);
        t = op(t, t);
        if (n % 2) {
            t = op(t, x);
        }
        return t;
    }
}
constexpr auto bpow(auto x, auto n, auto ans) {
    return bpow(x, n, ans, std::multiplies{});
}
template <typename T>
constexpr T bpow(T const &x, auto n) {
    return bpow(x, n, T(1));
}

// a * b mod m, without overflow
constexpr uint64_t mul_mod(uint64_t a, uint64_t b, uint64_t m) {
    return (uint64_t)((__uint128_t)a * b % m);
}
// x^n mod m, without overflow
constexpr uint64_t pow_mod(uint64_t x, uint64_t n, uint64_t m) {
    return bpow(x % m, n, 1 % m,
                [m](uint64_t a, uint64_t b) { return mul_mod(a, b, m); });
}

// Returns (g, x): g = gcd(a, n), x * a = g (mod n), 0 <= x < n / g.
// If g == 1, x is the inverse of a mod n.
constexpr std::pair<int64_t, int64_t> inv_gcd(int64_t a, int64_t n) {
    a = safe_mod(a, n);
    if (a == 0) return {n, 0};

    int64_t t = 0, newt = 1;
    int64_t r = n, newr = a;

    while (newr) {
        int64_t quotient = r / newr;
        r -= newr * quotient;
        t -= newt * quotient;

        std::swap(r, newr);
        std::swap(t, newt);
    }
    if (t < 0) t += n / r;
    return {r, t};
}

} // namespace algo::math
#line 4 "algo/math/combo.h"

namespace algo::math {

// Factorials and binomials mod p. Tables grow as needed.
//   combo<mint> c;
//   c.cmb(n, r);   // n choose r
//   c.perm(n, r);  // n! / (n - r)!
template <typename T>
struct combo {
    explicit combo(int n = 0) {
        if (n > 0) fact(n), inv_fact(n);
    }

    T fact(int n) {
        if (n >= (int)f.size()) {
            assert(n < mod());
            if (f.empty()) f.push_back(T(1));
            int m = grow_to(n, (int)f.size());
            f.reserve(m);
            for (int i = (int)f.size(); i < m; i++) {
                f.push_back(f.back() * T(i));
            }
        }
        return f[n];
    }
    T inv_fact(int n) {
        if (n >= (int)inv_f.size()) {
            assert(n < mod());
            if (inv_f.empty()) inv_f.push_back(T(1));
            int lo = (int)inv_f.size(), m = grow_to(n, lo);
            inv_f.resize(m);
            inv_f[m - 1] = T(1) / fact(m - 1);
            for (int i = m - 2; i >= lo; i--) {
                inv_f[i] = inv_f[i + 1] * T(i + 1);
            }
        }
        return inv_f[n];
    }
    T cmb(int n, int r) {
        if (r < 0 || r > n) {
            return T(0);
        } else {
            return fact(n) * inv_fact(r) * inv_fact(n - r);
        }
    }
    T perm(int n, int r) {
        if (r < 0 || r > n) {
            return T(0);
        } else {
            return fact(n) * inv_fact(n - r);
        }
    }

private:
    std::vector<T> f, inv_f;

    static int mod() {
        if constexpr (requires { T::mod(); }) {
            return T::mod();
        } else {
            return std::numeric_limits<int>::max();
        }
    }
    // New table size. Capped at mod, since n! = 0 for n >= mod.
    static int grow_to(int n, int cur) {
        return std::min<int64_t>(std::max<int64_t>(n + 1, 2LL * cur), mod());
    }
};

} // namespace algo::math

Back to top page