Codeforces #460 (Div. 2) E. Congruence Equation (R2100)

中国剰余定理のことをあれこれ調べていたら勢いで解いたん

問題へのリンク

問題概要

整数 が与えられる。

mod. )

が成立するような 以上 以下の整数 を数え上げよ。

制約

• 
• は素数
• <
• 

解法

Fermat の小定理から以下のことが言える:

• は mod. において周期 である
• は mod. において周期 である

したがって、 と は互いに素だから、 mod. は全体として周期 となることがわかる。

ここで、 の方を固定して集計することを考える。すなわち、 (mod. ) の場合について数えて条件を満たす の個数を数え上げて、それを について合算する。

(mod. )
⇔ (mod. )

となることから、すなわち

• (mod. )
• (mod. )

を満たす を数え上げることになる。中国剰余定理により、

(mod. )

を満たすような が求められるので、それを用いて 以下の整数を数え上げる。

#include <iostream>
#include <vector>
#include <cmath>
using namespace std;

inline long long mod(long long a, long long m) {
    return (a % m + m) % m;
}

long long inv(long long a, long long m) {
    long long b = m, u = 1, v = 0;
    while (b) {
        long long t = a/b;
        a -= t*b; swap(a, b);
        u -= t*v; swap(u, v);
    }
    return mod(u, m);
}

long long gcd(long long a, long long b) {
    if (b == 0) return a;
    else return gcd(b, a % b);
}

pair<long long, long long> ChineseRem(long long b1, long long m1, long long b2, long long m2) {   
        long long x = 0, m = 1;
        long long a = m, b = b1 - x, d = gcd(m1, a);
        if (b % d != 0) return make_pair(0, -1);
        long long t = mod(b / d * inv(a / d, m1 / d), m1 / d);
        x += m * t;
        m *= m1 / d;
        a = m, b = b2 - x, d = gcd(m2, a);
        if (b % d != 0) return make_pair(0, -1);
        t = mod(b / d * inv(a / d, m2 / d), m2 / d);
        x += m * t;
        m *= m2 / d;
        return make_pair(x % m, m);
}

long long a, b, p, x;

int main() {
  cin >> a >> b >> p >> x;
  long long res = 0;
  long long pow = 1;
  for (long long k = 0; k < p-1; ++k) {
    // ≡ k(mod. p-1), ≡ a^{-k}b (mod. p)
    long long b2 = inv(pow, p) * b % p;
    pair<long long, long long> c = ChineseRem(k, p-1, b2, p);
    long long tmp = (x + (c.second - c.first) % c.second) / c.second;
    res += tmp;
    pow = pow * a % p;
  }
  cout << res << endl;
}