Skip to content

extended_gcd.hpp

SECTIONMath INCLUDEnoya/math/extended_gcd.hpp

求扩展欧几里得系数、线性丢番图方程，并用 exCRT 合并多个模数不必互质的同余方程；用于不存在统一素数模前提的同余问题。

\[ \displaystyle ax+by=\gcd(a,b),\qquad x\equiv r_i\pmod{m_i}\quad(0\le i<k) \]

Complexity: Time: \(O(\log \max (\lvert a\rvert , \lvert b\rvert))\) for extgcd/crt, \(O(k \log M)\) for excrt, and \(O(k^{2} \log M)\) for Garner reconstruction. Space: \(O(\log \max (\lvert a\rvert , \lvert b\rvert))\) recursion, or \(O(k)\) for Garner reconstruction.

跳到代码 · GitHub ↗

Implementation

当前头文件，省略 include guard；依赖见 #include。

/// @complexity Time: O(log max(|a|,|b|)) for extgcd/crt,
/// O(k log M) for excrt, and O(k^2 log M) for Garner reconstruction.
/// Space: O(log max(|a|,|b|)) recursion, or O(k) for Garner reconstruction.

#include <algorithm>
#include <cassert>
#include <limits>
#include <numeric>
#include <vector>

namespace noya {

/// @brief Compute gcd(a, b) and coefficients x, y such that a*x + b*y = gcd.
/// T is a signed integer type; the nonnegative gcd must fit T.
template <typename T> T extgcd(T a, T b, T &x, T &y) {
  static_assert(std::numeric_limits<T>::is_integer &&
                std::numeric_limits<T>::is_signed);
  if (a == 0) {
    assert(b != std::numeric_limits<T>::lowest());
    x = 0;
    y = b < 0 ? -1 : 1;
    return b < 0 ? -b : b;
  }
  // In particular, avoid MIN / -1 even when the final coefficients fit.
  if (a == 1 || a == -1) {
    x = a;
    y = 0;
    return 1;
  }
  T p = b / a;
  T g = extgcd<T>(T(b % a), a, y, x);
  x -= p * y;
  return g;
}

/// @brief Solve a*x + b*y = c for signed integer inputs of at most 64 bits.
/// Returns false if no solution exists or gcd/coefficients cannot fit T.
template <typename T> bool diophantine(T a, T b, T c, T &x, T &y, T &g) {
  static_assert(std::numeric_limits<T>::is_integer &&
                std::numeric_limits<T>::is_signed &&
                std::numeric_limits<T>::digits <= 63);
  if (a == 0 && b == 0) {
    if (c == 0) {
      x = y = g = 0;
      return true;
    }
    return false;
  }
  const T lo = std::numeric_limits<T>::lowest();
  const T hi = std::numeric_limits<T>::max();
  if ((a == lo && (b == 0 || b == lo)) || (a == 0 && b == lo)) return false;
  T u, v;
  T div = extgcd(a, b, u, v);
  if (c % div != 0) return false;
  using i128 = __int128_t;
  i128 xx, yy;
  if (b == 0) {
    xx = i128(c) / a;
    yy = 0;
  } else {
    i128 mod = i128(b) / div;
    if (mod < 0) mod = -mod;
    xx = i128(u) * (i128(c) / div) % mod;
    // Center x to keep both coefficients small.
    if (xx > mod / 2) xx -= mod;
    if (xx < -mod / 2) xx += mod;
    yy = (i128(c) - i128(a) * xx) / b;
    if ((yy < lo || yy > hi) && a != 0) {
      // At signed endpoints, move to the nearest solution with y in range.
      i128 step = i128(a) / div;
      i128 dif = yy - (yy < lo ? lo : hi);
      i128 k = dif / step;
      if (dif % step != 0) k += (dif < 0) == (step < 0) ? 1 : -1;
      xx += k * (i128(b) / div);
      yy -= k * step;
    }
  }
  if (xx < lo || xx > hi || yy < lo || yy > hi) return false;
  x = T(xx);
  y = T(yy);
  g = div;
  return true;
}

/// @brief Chinese Remainder Theorem for two possibly non-coprime moduli.
/// Returns false if the congruences conflict or their lcm does not fit in a
/// signed 64-bit integer. Both moduli must be positive.
inline bool crt(long long k1, long long m1, long long k2, long long m2,
                long long &k, long long &m) {
  assert(m1 > 0 && m2 > 0);
  k1 %= m1;
  if (k1 < 0)
    k1 += m1;
  k2 %= m2;
  if (k2 < 0)
    k2 += m2;
  const long long g = std::gcd(m1, m2);
  using i128 = __int128_t;
  const i128 dif = i128(k2) - k1;
  if (dif % g != 0) {
    return false;
  }

  const i128 cm = i128(m1 / g) * m2;
  if (cm > std::numeric_limits<long long>::max()) {
    return false;
  }

  const long long rm = m2 / g;
  i128 ste = 0;
  if (rm != 1) {
    i128 inv = 0;
    i128 unu = 0;
    const i128 ig = extgcd(i128(m1 / g), i128(rm), inv, unu);
    assert(ig == 1);
    ste = dif / g % rm * inv % rm;
    if (ste < 0) {
      ste += rm;
    }
  }

  i128 cr = (i128(k1) + i128(m1) * ste) % cm;
  if (cr < 0) {
    cr += cm;
  }
  k = static_cast<long long>(cr);
  m = static_cast<long long>(cm);
  return true;
}

/// @brief Extended CRT for two possibly non-coprime congruences.
inline bool excrt(long long k1, long long m1, long long k2, long long m2,
                  long long &k, long long &m) {
  return crt(k1, m1, k2, m2, k, m);
}

/// @brief Extended CRT for x = rsd[i] (mod ms[i]). The returned
/// residue lies in [0, mod); the empty system returns (0, 1).
inline bool excrt(const std::vector<long long> &rsd,
                  const std::vector<long long> &ms, long long &re0,
                  long long &mod) {
  assert(rsd.size() == ms.size());
  long long mr = 0;
  long long mm = 1;
  for (std::size_t i = 0; i < rsd.size(); i++) {
    assert(ms[i] > 0);
    long long nr = 0;
    long long nm = 0;
    if (!crt(mr, mm, rsd[i], ms[i], nr, nm)) {
      return false;
    }
    mr = nr;
    mm = nm;
  }
  re0 = mr;
  mod = mm;
  return true;
}

/// @brief Garner's algorithm for positive pairwise-coprime moduli; stores the
/// least nonnegative result in res. The product of the moduli must fit T.
template <typename T>
void crt_garner(const std::vector<int> &p, const std::vector<int> &a, T &res) {
  assert(p.size() == a.size());
  auto inv = [&](int q, int m) {
    q %= m;
    if (q < 0)
      q += m;
    int b = m, u = 0, v = 1;
    while (q) {
      int t = b / q;
      b -= t * q;
      std::swap(q, b);
      u -= t * v;
      std::swap(u, v);
    }
    assert(b == 1);
    if (u < 0)
      u += m;
    return u;
  };
  std::vector<int> x(p.size());
  for (int i = 0; i < (int)p.size(); i++) {
    assert(0 <= a[i] && a[i] < p[i]);
    x[i] = a[i];
    for (int j = 0; j < i; j++) {
      x[i] = (int)((long long)(x[i] - x[j]) * inv(p[j], p[i]) % p[i]);
      if (x[i] < 0)
        x[i] += p[i];
    }
  }
  res = 0;
  for (int i = (int)p.size() - 1; i >= 0; i--) {
    res = res * p[i] + x[i];
  }
}

} // namespace noya

 // NOYA_EXTENDED_GCD_HPP

#ifndef NOYA_EXTENDED_GCD_HPP
#define NOYA_EXTENDED_GCD_HPP 1

/// @complexity Time: O(log max(|a|,|b|)) for extgcd/crt,
/// O(k log M) for excrt, and O(k^2 log M) for Garner reconstruction.
/// Space: O(log max(|a|,|b|)) recursion, or O(k) for Garner reconstruction.

#include <algorithm>
#include <cassert>
#include <limits>
#include <numeric>
#include <vector>

namespace noya {

/// @brief Compute gcd(a, b) and coefficients x, y such that a*x + b*y = gcd.
/// T is a signed integer type; the nonnegative gcd must fit T.
template <typename T> T extgcd(T a, T b, T &x, T &y) {
  static_assert(std::numeric_limits<T>::is_integer &&
                std::numeric_limits<T>::is_signed);
  if (a == 0) {
    assert(b != std::numeric_limits<T>::lowest());
    x = 0;
    y = b < 0 ? -1 : 1;
    return b < 0 ? -b : b;
  }
  // In particular, avoid MIN / -1 even when the final coefficients fit.
  if (a == 1 || a == -1) {
    x = a;
    y = 0;
    return 1;
  }
  T p = b / a;
  T g = extgcd<T>(T(b % a), a, y, x);
  x -= p * y;
  return g;
}

/// @brief Solve a*x + b*y = c for signed integer inputs of at most 64 bits.
/// Returns false if no solution exists or gcd/coefficients cannot fit T.
template <typename T> bool diophantine(T a, T b, T c, T &x, T &y, T &g) {
  static_assert(std::numeric_limits<T>::is_integer &&
                std::numeric_limits<T>::is_signed &&
                std::numeric_limits<T>::digits <= 63);
  if (a == 0 && b == 0) {
    if (c == 0) {
      x = y = g = 0;
      return true;
    }
    return false;
  }
  const T lo = std::numeric_limits<T>::lowest();
  const T hi = std::numeric_limits<T>::max();
  if ((a == lo && (b == 0 || b == lo)) || (a == 0 && b == lo)) return false;
  T u, v;
  T div = extgcd(a, b, u, v);
  if (c % div != 0) return false;
  using i128 = __int128_t;
  i128 xx, yy;
  if (b == 0) {
    xx = i128(c) / a;
    yy = 0;
  } else {
    i128 mod = i128(b) / div;
    if (mod < 0) mod = -mod;
    xx = i128(u) * (i128(c) / div) % mod;
    // Center x to keep both coefficients small.
    if (xx > mod / 2) xx -= mod;
    if (xx < -mod / 2) xx += mod;
    yy = (i128(c) - i128(a) * xx) / b;
    if ((yy < lo || yy > hi) && a != 0) {
      // At signed endpoints, move to the nearest solution with y in range.
      i128 step = i128(a) / div;
      i128 dif = yy - (yy < lo ? lo : hi);
      i128 k = dif / step;
      if (dif % step != 0) k += (dif < 0) == (step < 0) ? 1 : -1;
      xx += k * (i128(b) / div);
      yy -= k * step;
    }
  }
  if (xx < lo || xx > hi || yy < lo || yy > hi) return false;
  x = T(xx);
  y = T(yy);
  g = div;
  return true;
}

/// @brief Chinese Remainder Theorem for two possibly non-coprime moduli.
/// Returns false if the congruences conflict or their lcm does not fit in a
/// signed 64-bit integer. Both moduli must be positive.
inline bool crt(long long k1, long long m1, long long k2, long long m2,
                long long &k, long long &m) {
  assert(m1 > 0 && m2 > 0);
  k1 %= m1;
  if (k1 < 0)
    k1 += m1;
  k2 %= m2;
  if (k2 < 0)
    k2 += m2;
  const long long g = std::gcd(m1, m2);
  using i128 = __int128_t;
  const i128 dif = i128(k2) - k1;
  if (dif % g != 0) {
    return false;
  }

  const i128 cm = i128(m1 / g) * m2;
  if (cm > std::numeric_limits<long long>::max()) {
    return false;
  }

  const long long rm = m2 / g;
  i128 ste = 0;
  if (rm != 1) {
    i128 inv = 0;
    i128 unu = 0;
    const i128 ig = extgcd(i128(m1 / g), i128(rm), inv, unu);
    assert(ig == 1);
    ste = dif / g % rm * inv % rm;
    if (ste < 0) {
      ste += rm;
    }
  }

  i128 cr = (i128(k1) + i128(m1) * ste) % cm;
  if (cr < 0) {
    cr += cm;
  }
  k = static_cast<long long>(cr);
  m = static_cast<long long>(cm);
  return true;
}

/// @brief Extended CRT for two possibly non-coprime congruences.
inline bool excrt(long long k1, long long m1, long long k2, long long m2,
                  long long &k, long long &m) {
  return crt(k1, m1, k2, m2, k, m);
}

/// @brief Extended CRT for x = rsd[i] (mod ms[i]). The returned
/// residue lies in [0, mod); the empty system returns (0, 1).
inline bool excrt(const std::vector<long long> &rsd,
                  const std::vector<long long> &ms, long long &re0,
                  long long &mod) {
  assert(rsd.size() == ms.size());
  long long mr = 0;
  long long mm = 1;
  for (std::size_t i = 0; i < rsd.size(); i++) {
    assert(ms[i] > 0);
    long long nr = 0;
    long long nm = 0;
    if (!crt(mr, mm, rsd[i], ms[i], nr, nm)) {
      return false;
    }
    mr = nr;
    mm = nm;
  }
  re0 = mr;
  mod = mm;
  return true;
}

/// @brief Garner's algorithm for positive pairwise-coprime moduli; stores the
/// least nonnegative result in res. The product of the moduli must fit T.
template <typename T>
void crt_garner(const std::vector<int> &p, const std::vector<int> &a, T &res) {
  assert(p.size() == a.size());
  auto inv = [&](int q, int m) {
    q %= m;
    if (q < 0)
      q += m;
    int b = m, u = 0, v = 1;
    while (q) {
      int t = b / q;
      b -= t * q;
      std::swap(q, b);
      u -= t * v;
      std::swap(u, v);
    }
    assert(b == 1);
    if (u < 0)
      u += m;
    return u;
  };
  std::vector<int> x(p.size());
  for (int i = 0; i < (int)p.size(); i++) {
    assert(0 <= a[i] && a[i] < p[i]);
    x[i] = a[i];
    for (int j = 0; j < i; j++) {
      x[i] = (int)((long long)(x[i] - x[j]) * inv(p[j], p[i]) % p[i]);
      if (x[i] < 0)
        x[i] += p[i];
    }
  }
  res = 0;
  for (int i = (int)p.size() - 1; i >= 0; i--) {
    res = res * p[i] + x[i];
  }
}

} // namespace noya

#endif // NOYA_EXTENDED_GCD_HPP

/// @complexity Time: O(log max(|a|,|b|)) for extgcd/crt,
/// O(k log M) for excrt, and O(k^2 log M) for Garner reconstruction.
/// Space: O(log max(|a|,|b|)) recursion, or O(k) for Garner reconstruction.

#include <algorithm>
#include <cassert>
#include <limits>
#include <numeric>
#include <vector>

namespace noya {

/// @brief Compute gcd(a, b) and coefficients x, y such that a*x + b*y = gcd.
/// T is a signed integer type; the nonnegative gcd must fit T.
template <typename T> T extgcd(T a, T b, T &x, T &y) {
  static_assert(std::numeric_limits<T>::is_integer &&
                std::numeric_limits<T>::is_signed);
  if (a == 0) {
    assert(b != std::numeric_limits<T>::lowest());
    x = 0;
    y = b < 0 ? -1 : 1;
    return b < 0 ? -b : b;
  }
  // In particular, avoid MIN / -1 even when the final coefficients fit.
  if (a == 1 || a == -1) {
    x = a;
    y = 0;
    return 1;
  }
  T p = b / a;
  T g = extgcd<T>(T(b % a), a, y, x);
  x -= p * y;
  return g;
}

/// @brief Solve a*x + b*y = c for signed integer inputs of at most 64 bits.
/// Returns false if no solution exists or gcd/coefficients cannot fit T.
template <typename T> bool diophantine(T a, T b, T c, T &x, T &y, T &g) {
  static_assert(std::numeric_limits<T>::is_integer &&
                std::numeric_limits<T>::is_signed &&
                std::numeric_limits<T>::digits <= 63);
  if (a == 0 && b == 0) {
    if (c == 0) {
      x = y = g = 0;
      return true;
    }
    return false;
  }
  const T lo = std::numeric_limits<T>::lowest();
  const T hi = std::numeric_limits<T>::max();
  if ((a == lo && (b == 0 || b == lo)) || (a == 0 && b == lo)) return false;
  T u, v;
  T div = extgcd(a, b, u, v);
  if (c % div != 0) return false;
  using i128 = __int128_t;
  i128 xx, yy;
  if (b == 0) {
    xx = i128(c) / a;
    yy = 0;
  } else {
    i128 mod = i128(b) / div;
    if (mod < 0) mod = -mod;
    xx = i128(u) * (i128(c) / div) % mod;
    // Center x to keep both coefficients small.
    if (xx > mod / 2) xx -= mod;
    if (xx < -mod / 2) xx += mod;
    yy = (i128(c) - i128(a) * xx) / b;
    if ((yy < lo || yy > hi) && a != 0) {
      // At signed endpoints, move to the nearest solution with y in range.
      i128 step = i128(a) / div;
      i128 dif = yy - (yy < lo ? lo : hi);
      i128 k = dif / step;
      if (dif % step != 0) k += (dif < 0) == (step < 0) ? 1 : -1;
      xx += k * (i128(b) / div);
      yy -= k * step;
    }
  }
  if (xx < lo || xx > hi || yy < lo || yy > hi) return false;
  x = T(xx);
  y = T(yy);
  g = div;
  return true;
}

/// @brief Chinese Remainder Theorem for two possibly non-coprime moduli.
/// Returns false if the congruences conflict or their lcm does not fit in a
/// signed 64-bit integer. Both moduli must be positive.
inline bool crt(long long k1, long long m1, long long k2, long long m2,
                long long &k, long long &m) {
  assert(m1 > 0 && m2 > 0);
  k1 %= m1;
  if (k1 < 0)
    k1 += m1;
  k2 %= m2;
  if (k2 < 0)
    k2 += m2;
  const long long g = std::gcd(m1, m2);
  using i128 = __int128_t;
  const i128 dif = i128(k2) - k1;
  if (dif % g != 0) {
    return false;
  }

  const i128 cm = i128(m1 / g) * m2;
  if (cm > std::numeric_limits<long long>::max()) {
    return false;
  }

  const long long rm = m2 / g;
  i128 ste = 0;
  if (rm != 1) {
    i128 inv = 0;
    i128 unu = 0;
    const i128 ig = extgcd(i128(m1 / g), i128(rm), inv, unu);
    assert(ig == 1);
    ste = dif / g % rm * inv % rm;
    if (ste < 0) {
      ste += rm;
    }
  }

  i128 cr = (i128(k1) + i128(m1) * ste) % cm;
  if (cr < 0) {
    cr += cm;
  }
  k = static_cast<long long>(cr);
  m = static_cast<long long>(cm);
  return true;
}

/// @brief Extended CRT for two possibly non-coprime congruences.
inline bool excrt(long long k1, long long m1, long long k2, long long m2,
                  long long &k, long long &m) {
  return crt(k1, m1, k2, m2, k, m);
}

/// @brief Extended CRT for x = rsd[i] (mod ms[i]). The returned
/// residue lies in [0, mod); the empty system returns (0, 1).
inline bool excrt(const std::vector<long long> &rsd,
                  const std::vector<long long> &ms, long long &re0,
                  long long &mod) {
  assert(rsd.size() == ms.size());
  long long mr = 0;
  long long mm = 1;
  for (std::size_t i = 0; i < rsd.size(); i++) {
    assert(ms[i] > 0);
    long long nr = 0;
    long long nm = 0;
    if (!crt(mr, mm, rsd[i], ms[i], nr, nm)) {
      return false;
    }
    mr = nr;
    mm = nm;
  }
  re0 = mr;
  mod = mm;
  return true;
}

/// @brief Garner's algorithm for positive pairwise-coprime moduli; stores the
/// least nonnegative result in res. The product of the moduli must fit T.
template <typename T>
void crt_garner(const std::vector<int> &p, const std::vector<int> &a, T &res) {
  assert(p.size() == a.size());
  auto inv = [&](int q, int m) {
    q %= m;
    if (q < 0)
      q += m;
    int b = m, u = 0, v = 1;
    while (q) {
      int t = b / q;
      b -= t * q;
      std::swap(q, b);
      u -= t * v;
      std::swap(u, v);
    }
    assert(b == 1);
    if (u < 0)
      u += m;
    return u;
  };
  std::vector<int> x(p.size());
  for (int i = 0; i < (int)p.size(); i++) {
    assert(0 <= a[i] && a[i] < p[i]);
    x[i] = a[i];
    for (int j = 0; j < i; j++) {
      x[i] = (int)((long long)(x[i] - x[j]) * inv(p[j], p[i]) % p[i]);
      if (x[i] < 0)
        x[i] += p[i];
    }
  }
  res = 0;
  for (int i = (int)p.size() - 1; i >= 0; i--) {
    res = res * p[i] + x[i];
  }
}

} // namespace noya

 // NOYA_EXTENDED_GCD_HPP