Mod/mod_pow.hpp

• View this file on GitHub
• Last update: 2024-05-29 22:35:44+07:00
• Include: #include "Mod/mod_pow.hpp"

Depends on

• Modint/Barrett_reduction.hpp
• Modint/dynamic_modint.hpp

Required by

• Mod/mod_sqrt.hpp

Verified with

• test/yosupo/Math/mod_sqrt_yosupo.test.cpp

Code

#pragma once
#include "Modint/dynamic_modint.hpp"
u32 mod_pow(int a, ll n, int mod) {
  assert(n >= 0);
  a = ((a %= mod) < 0 ? a + mod : a);
  if ((mod & 1) && (mod < (1 << 30))) {
    using mint = dynamic_modint<202311021>;
    mint::set_mod(mod);
    return mint(a).pow(n).get();
  }
  Barrett bt(mod);
  int r = 1;
  while (n) {
    if (n & 1) r = bt.mul(r, a);
    a = bt.mul(a, a), n >>= 1;
  }
  return r;
}

u64 mod_pow_64(ll a, ll n, u64 mod) {
  assert(n >= 0);
  a = ((a %= mod) < 0 ? a + mod : a);
  if ((mod & 1) && (mod < (u64(1) << 62))) {
    using mint =dynamic_modint<202311021>;
    mint::set_mod(mod);
    return mint(a).pow(n).get();
  }
  Barrett_64 bt(mod);
  ll r = 1;
  while (n) {
    if (n & 1) r = bt.mul(r, a);
    a = bt.mul(a, a), n >>= 1;
  }
  return r;
}

#line 2 "Modint/Barrett_reduction.hpp"
/*
  @see https://nyaannyaan.github.io/library/modint/barrett-reduction.hpp
  @see https://en.wikipedia.org/wiki/Barrett_reduction
*/
struct Barrett {
  using u32 = unsigned int;
  using i64 = long long;
  using u64 = unsigned long long;
  u32 m;
  u64 im;
  Barrett() : m(), im() {}
  Barrett(int n) : m(n), im(u64(-1) / m + 1) {}
  constexpr inline i64 quo(u64 n) {
    u64 x = u64((__uint128_t(n) * im) >> 64);
    u32 r = n - x * m;
    return m <= r ? x - 1 : x;
  }
  constexpr inline i64 rem(u64 n) {
    u64 x = u64((__uint128_t(n) * im) >> 64);
    u32 r = n - x * m;
    return m <= r ? r + m : r;
  }
  constexpr inline pair<i64, int> quorem(u64 n) {
    u64 x = u64((__uint128_t(n) * im) >> 64);
    u32 r = n - x * m;
    if (m <= r) return {x - 1, r + m};
    return {x, r};
  }
  constexpr inline i64 pow(u64 n, i64 p) {
    u32 a = rem(n), r = m == 1 ? 0 : 1;
    while (p) {
      if (p & 1) r = rem(u64(r) * a);
      a = rem(u64(a) * a);
      p >>= 1;
    }
    return r;
  }
  constexpr inline u32 mul(u32 a, u32 b) {
    return rem(u64(a) * b);
  }
};

//u64 version:
struct Barrett_64 {
  u128 mod, mh, ml;

  explicit Barrett_64(u64 mod = 1) : mod(mod) {
    u128 m = u128(-1) / mod;
    if (m * mod + mod == u128(0)) ++m;
    mh = m >> 64;
    ml = m & u64(-1);
  }

  u64 umod() const { return mod; }

  u64 rem(u128 x) {
    u128 z = (x & u64(-1)) * ml;
    z = (x & u64(-1)) * mh + (x >> 64) * ml + (z >> 64);
    z = (x >> 64) * mh + (z >> 64);
    x -= z * mod;
    return x < mod ? x : x - mod;
  }

  u64 mul(u64 a, u64 b) { return rem(u128(a) * b); }
};
#line 3 "Modint/dynamic_modint.hpp"
template <int id>
struct dynamic_modint {
  int x;

  dynamic_modint() : x(0) {}

  dynamic_modint(int64_t y) {
    int z = y % get_mod();
    if (z < 0) z += get_mod();
    x = z;
  }

  dynamic_modint &operator+=(const dynamic_modint &p) {
    if ((x += p.x) >= get_mod()) x -= get_mod();
    return *this;
  }

  dynamic_modint &operator-=(const dynamic_modint &p) {
    if ((x += get_mod() - p.x) >= get_mod()) x -= get_mod();
    return *this;
  }

  dynamic_modint &operator*=(const dynamic_modint &p) {
    x = rem((unsigned long long)x * p.x);
    return *this;
  }

  dynamic_modint &operator/=(const dynamic_modint &p) {
    *this *= p.inv();
    return *this;
  }

  dynamic_modint operator-() const { return dynamic_modint(-x); }
  dynamic_modint operator+() const { return *this; }

  dynamic_modint operator+(const dynamic_modint &p) const {
    return dynamic_modint(*this) += p;
  }

  dynamic_modint operator-(const dynamic_modint &p) const {
    return dynamic_modint(*this) -= p;
  }

  dynamic_modint operator*(const dynamic_modint &p) const {
    return dynamic_modint(*this) *= p;
  }

  dynamic_modint operator/(const dynamic_modint &p) const {
    return dynamic_modint(*this) /= p;
  }

  bool operator==(const dynamic_modint &p) const { return x == p.x; }

  bool operator!=(const dynamic_modint &p) const { return x != p.x; }

  dynamic_modint inv() const {
    int a = x, b = get_mod(), u = 1, v = 0, t;
    while (b > 0) {
      t = a / b;
      swap(a -= t * b, b);
      swap(u -= t * v, v);
    }
    return dynamic_modint(u);
  }

  dynamic_modint pow(int64_t n) const {
    dynamic_modint ret(1), mul(x);
    while (n > 0) {
      if (n & 1) ret *= mul;
      mul *= mul;
      n >>= 1;
    }
    return ret;
  }

  friend ostream &operator<<(ostream &os, const dynamic_modint &p) {
    return os << p.x;
  }

  friend istream &operator>>(istream &is, dynamic_modint &a) {
    int64_t t;
    is >> t;
    a = dynamic_modint(t);
    return (is);
  }

  int get() const { return x; }

  inline unsigned int rem(unsigned long long p) { return barrett().rem(p); }

  static inline Barrett &barrett() {
    static Barrett b;
    return b;
  }

  static inline int &get_mod() {
    static int mod = 0;
    return mod;
  }

  static void set_mod(int md) {
    assert(0 < md && md <= (1LL << 30) - 1);
    get_mod() = md;
    barrett() = Barrett(md);
  }
};

using modint = dynamic_modint<-1>;
#line 3 "Mod/mod_pow.hpp"
u32 mod_pow(int a, ll n, int mod) {
  assert(n >= 0);
  a = ((a %= mod) < 0 ? a + mod : a);
  if ((mod & 1) && (mod < (1 << 30))) {
    using mint = dynamic_modint<202311021>;
    mint::set_mod(mod);
    return mint(a).pow(n).get();
  }
  Barrett bt(mod);
  int r = 1;
  while (n) {
    if (n & 1) r = bt.mul(r, a);
    a = bt.mul(a, a), n >>= 1;
  }
  return r;
}

u64 mod_pow_64(ll a, ll n, u64 mod) {
  assert(n >= 0);
  a = ((a %= mod) < 0 ? a + mod : a);
  if ((mod & 1) && (mod < (u64(1) << 62))) {
    using mint =dynamic_modint<202311021>;
    mint::set_mod(mod);
    return mint(a).pow(n).get();
  }
  Barrett_64 bt(mod);
  ll r = 1;
  while (n) {
    if (n & 1) r = bt.mul(r, a);
    a = bt.mul(a, a), n >>= 1;
  }
  return r;
}

Back to top page