Library/Math/Math.hpp

• View this file on GitHub
• Last update: 2025-01-02 21:02:45+09:00
• Include: #include "Library/Math/Math.hpp"

Depends on

• Library/Math/EuclideanAlgorithm.hpp

Required by

• Library/Math/DynamicModInt.hpp
• Library/Math/ModInt.hpp

Verified with

• Test/DataStructure/DisjointSetUnion_Potential.test.cpp
• Test/Graph/Tree/ReRootingDP.test.cpp
• Test/Math/Convolution_and.test.cpp
• Test/Math/Math.test.cpp
• Test/Math/Math_dynamic.test.cpp
• Test/Math/Math_log.test.cpp
• Test/Math/Math_pow.test.cpp
• Test/Math/Math_pow_dynamic.test.cpp
• Test/Math/Matrix_pow.test.cpp

Code

#pragma once

#include <cmath>

#include <numeric>

#include <optional>

#include <ranges>

#include <unordered_map>

#include <vector>


#include "./EuclideanAlgorithm.hpp"


namespace mtd {
  template <class T>
  class Math {
    const std::vector<T> m_fac;
    const std::vector<T> m_finv;

    auto constructFac(int s) {
      std::vector<T> fac(s);
      fac[0] = fac[1] = 1;
      for (long long i = 2; i < s; ++i) { fac[i] = fac[i - 1] * i; }
      return fac;
    }
    auto constructInv(int s) {
      std::vector<T> finv(s);
      finv[s - 1] = 1 / m_fac[s - 1];
      for (long long i = s - 2; i >= 0; --i) {
        finv[i] = finv[i + 1] * (i + 1);
      }
      return finv;
    }

  public:
    constexpr Math(int size = 3 * static_cast<int>(1e6))
        : m_fac(constructFac(size)), m_finv(constructInv(size)) {}

    /* O(log b) */
    static constexpr T pow(T a, long long b) {
      T ans = 1;
      while (b > 0) {
        if (b & 1) { ans *= a; }
        b >>= 1;
        a *= a;
      }
      return ans;
    }

    /* O(log mod) */
    template <class S>
    static constexpr std::optional<long long> log(S x, S y, S mod) {
      x %= mod;
      y %= mod;

      if (mod == 1) { return 0; }
      if (x == 0 && y == 0) { return 1; }
      if (x == 0 && y == 1) { return 0; }
      if (x == 0) { return std::nullopt; }
      if (y == 1) { return 0; }

      if (auto g = std::gcd(x, mod); g > 1) {
        if (y % g) { return std::nullopt; }
        auto nx = x / g;
        auto nmod = mod / g;
        auto ea = mtd::EuclideanAlgorithm(nx, -nmod, 1);
        auto [t, _] = ea.getMinX();
        auto [nx_inv, __] = ea.get(t);
        nx_inv %= nmod;
        if (auto ans = log(x, y / g * nx_inv, nmod); ans) {
          return ans.value() + 1;
        } else {
          return ans;
        }
      }

      auto s = static_cast<S>(std::sqrt(mod));
      S xe = y;
      std::unordered_map<S, S> map;
      map.reserve(s);
      for (auto i : std::views::iota(0, s)) {
        (xe *= x) %= mod;
        map[xe] = i + 1;
      }

      S xs = 1;
      for ([[maybe_unused]] auto _ : std::views::iota(0, s)) {
        (xs *= x) %= mod;
      }
      S xse = 1;
      for (auto i : std::views::iota(0, mod / s + 5)) {
        (xse *= xs) %= mod;
        if (map.contains(xse)) { return s * (i + 1) - map[xse]; }
      }
      return std::nullopt;
    }
    constexpr std::optional<long long> log(long long x,
                                           long long y) requires requires {
      typename T::value_type;
      T::mod();
    }
    { return log(x, y, T::mod()); }

    constexpr auto fact(int n) const { return (n < 0) ? 0 : m_fac[n]; }
    constexpr auto factInv(int n) const { return (n < 0 ? 0 : m_finv[n]); }
    constexpr auto comb(int n, int r) const {
      return fact(n) * factInv(r) * factInv(n - r);
    }
    constexpr auto perm(int n, int r) const { return fact(n) * factInv(n - r); }
  };
}  // namespace mtd

#line 2 "Library/Math/Math.hpp"

#include <cmath>

#include <numeric>

#include <optional>

#include <ranges>

#include <unordered_map>

#include <vector>


#line 2 "Library/Math/EuclideanAlgorithm.hpp"

#include <iostream>

#line 5 "Library/Math/EuclideanAlgorithm.hpp"
#include <tuple>


namespace mtd {

  class EuclideanAlgorithm {
    using T = long long;

    // 大きすぎるとオーバーフローしてしまう

    const static inline T m_mx = 1e9;

    const T m_a;
    const T m_b;
    const T m_c;

    T m_gcd;
    T m_x;
    T m_y;

    auto excludedEuclidAlgorithm(T a, T b) -> std::tuple<T, T, T> {
      if (a < 0) {
        auto [g, x, y] = excludedEuclidAlgorithm(-a, -b);
        return {g, -x, -y};
      }
      if (b == 0) { return {a, 1, 0}; }
      auto [g, y, x] = excludedEuclidAlgorithm(b, a % b);
      y -= a / b * x;
      return {g, x, y};
    }

    auto kRange(T x, T b, T l) const -> std::pair<T, T> {
      // x + b * k >= l を満たす k の範囲を求める

      T xd = (l - x);
      if (b == 0 && x >= l) { return {-m_mx, m_mx}; }
      if (b == 0 && x < l) { return {m_mx, -m_mx}; }
      if (b > 0 && xd < 0) { return {xd / b, m_mx}; }
      if (b > 0 && xd >= 0) { return {(xd + b - 1) / b, m_mx}; }
      if (b < 0 && xd < 0) { return {-m_mx, (-xd) / (-b)}; }
      if (b < 0 && xd >= 0) { return {-m_mx, -(xd - b - 1) / (-b)}; }
      return {m_mx, -m_mx};
    }

  public:
    auto debug() const {
      std::cout << m_a << " * " << m_x << " + " << m_b << " * " << m_y << " = "
                << m_c << std::endl;
      std::cout << "calc: " << m_a * m_x + m_b * m_y << " = " << m_c
                << std::endl;
    }

    EuclideanAlgorithm(T a, T b, T c) : m_a(a), m_b(b), m_c(c) {
      if (a == 0 && b == 0) { throw std::runtime_error(""); }
      auto [g, x, y] = excludedEuclidAlgorithm(a, b);
      if (c % g > 0) {
        throw std::runtime_error(
            "There is no solution to the equation. c must be divisible by "
            "gcd(a,b).");
      }
      m_gcd = g;
      m_x = c / g * x;
      m_y = c / g * y;
    }
    EuclideanAlgorithm(T a, T b) : EuclideanAlgorithm(a, b, std::gcd(a, b)) {}

    auto gcd() const { return m_gcd; }
    auto get(T x, T y) const { return m_a * x + m_b * y; }
    auto get(T k) const -> std::pair<T, T> {
      if (m_b == 0) { return {m_x, m_y - k}; }
      if (m_a == 0) { return {m_x + k, m_y}; }
      return {m_x + m_b * k, m_y - m_a * k};
    }
    // x>=x_lとなるようなkの範囲

    auto getMinX(T x_l = 0) const -> std::pair<T, T> {
      return kRange(m_x, m_b, x_l);
    }
    // y>=y_lとなるようなkの範囲

    auto getMinY(T y_l = 0) const -> std::pair<T, T> {
      return kRange(m_y, -1 * m_a, y_l);
    }
    // x>=x_l, y>=y_lとなるようなkの範囲

    auto getMin(T x_l = 0, T y_l = 0) const -> std::pair<T, T> {
      auto [xl, xr] = getMinX(x_l);
      auto [yl, yr] = getMinY(y_l);
      return {std::max(xl, yl), std::min(xr, yr)};
    }
  };

}  // namespace mtd

#line 11 "Library/Math/Math.hpp"

namespace mtd {
  template <class T>
  class Math {
    const std::vector<T> m_fac;
    const std::vector<T> m_finv;

    auto constructFac(int s) {
      std::vector<T> fac(s);
      fac[0] = fac[1] = 1;
      for (long long i = 2; i < s; ++i) { fac[i] = fac[i - 1] * i; }
      return fac;
    }
    auto constructInv(int s) {
      std::vector<T> finv(s);
      finv[s - 1] = 1 / m_fac[s - 1];
      for (long long i = s - 2; i >= 0; --i) {
        finv[i] = finv[i + 1] * (i + 1);
      }
      return finv;
    }

  public:
    constexpr Math(int size = 3 * static_cast<int>(1e6))
        : m_fac(constructFac(size)), m_finv(constructInv(size)) {}

    /* O(log b) */
    static constexpr T pow(T a, long long b) {
      T ans = 1;
      while (b > 0) {
        if (b & 1) { ans *= a; }
        b >>= 1;
        a *= a;
      }
      return ans;
    }

    /* O(log mod) */
    template <class S>
    static constexpr std::optional<long long> log(S x, S y, S mod) {
      x %= mod;
      y %= mod;

      if (mod == 1) { return 0; }
      if (x == 0 && y == 0) { return 1; }
      if (x == 0 && y == 1) { return 0; }
      if (x == 0) { return std::nullopt; }
      if (y == 1) { return 0; }

      if (auto g = std::gcd(x, mod); g > 1) {
        if (y % g) { return std::nullopt; }
        auto nx = x / g;
        auto nmod = mod / g;
        auto ea = mtd::EuclideanAlgorithm(nx, -nmod, 1);
        auto [t, _] = ea.getMinX();
        auto [nx_inv, __] = ea.get(t);
        nx_inv %= nmod;
        if (auto ans = log(x, y / g * nx_inv, nmod); ans) {
          return ans.value() + 1;
        } else {
          return ans;
        }
      }

      auto s = static_cast<S>(std::sqrt(mod));
      S xe = y;
      std::unordered_map<S, S> map;
      map.reserve(s);
      for (auto i : std::views::iota(0, s)) {
        (xe *= x) %= mod;
        map[xe] = i + 1;
      }

      S xs = 1;
      for ([[maybe_unused]] auto _ : std::views::iota(0, s)) {
        (xs *= x) %= mod;
      }
      S xse = 1;
      for (auto i : std::views::iota(0, mod / s + 5)) {
        (xse *= xs) %= mod;
        if (map.contains(xse)) { return s * (i + 1) - map[xse]; }
      }
      return std::nullopt;
    }
    constexpr std::optional<long long> log(long long x,
                                           long long y) requires requires {
      typename T::value_type;
      T::mod();
    }
    { return log(x, y, T::mod()); }

    constexpr auto fact(int n) const { return (n < 0) ? 0 : m_fac[n]; }
    constexpr auto factInv(int n) const { return (n < 0 ? 0 : m_finv[n]); }
    constexpr auto comb(int n, int r) const {
      return fact(n) * factInv(r) * factInv(n - r);
    }
    constexpr auto perm(int n, int r) const { return fact(n) * factInv(n - r); }
  };
}  // namespace mtd

Back to top page