Yukicoder No.551 夏休みの思い出(2)

,

https://yukicoder.me/problems/no/551

CTF感がすごい

problem

一般の$2$次合同式$ax^2 + bx + c \equiv 0 \pmod{p}$を解く問題。

solution

有限体$\mathbb{F}_p$の上では$2$次方程式の解の公式が使える。 判別式$D = b^2 - 4ac$とおいて$w^2 \equiv D \pmod{p}$なる$w$が存在すれば$x \equiv (- b \pm w)(2a)^{-1} \pmod{p}$が解。存在しなければ解なし。

$\mathbb{F}_p$で平方根を計算する問題に帰着された。 平方剰余だとかLegendre記号だとかEuler規準だとかそういうので検索すれば解ける。 原始根が与えられているが使わない。 計算量について、$O(\sqrt{P} + Q \log P)$ぐらいで解けそうだが正確なのは分からない。

離散対数の特殊な場合なのでbaby-step giant-stepでえいってしてもできたらしい。それはそうだが気付かなかった。

implementation

#include <algorithm>
#include <cassert>
#include <cstdio>
#include <vector>
#define repeat(i, n) for (int i = 0; (i) < int(n); ++(i))
#define whole(f, x, ...) ([&](decltype((x)) whole) { return (f)(begin(whole), end(whole), ## __VA_ARGS__); })(x)
using ll = long long;
using namespace std;

ll powmod(ll x, ll y, ll p) { // O(log y)
    assert (0 <= x and x < p);
    assert (0 <= y);
    ll z = 1;
    for (ll i = 1; i <= y; i <<= 1) {
        if (y & i) z = z * x % p;
        x = x * x % p;
    }
    return z;
}
ll modinv(ll x, ll p) { // p must be a prime, O(log p)
    assert ((x % p + p) % p != 0);
    return powmod(x, p - 2, p);
}

vector<bool> sieve_of_eratosthenes(int n) { // enumerate primes in [2,n] with O(n log log n)
    vector<bool> is_prime(n+1, true);
    is_prime[0] = is_prime[1] = false;
    for (int i = 2; i*i <= n; ++i)
        if (is_prime[i])
            for (int k = i+i; k <= n; k += i)
                is_prime[k] = false;
    return is_prime;
}
vector<int> list_primes(int n) {
    auto is_prime = sieve_of_eratosthenes(n);
    vector<int> primes;
    for (int i = 2; i <= n; ++i)
        if (is_prime[i])
            primes.push_back(i);
    return primes;
}

int legendre_symbol(int a, int p) {
    return powmod(a, (p - 1) / 2, p); // Euler's criterion
}
// transported from https://github.com/koba-e964/lib-number/blob/master/modsqrt.rb
int modsqrt(int a, int p) {
    a %= p;
    if (a == 0) return 0;
    if (p == 2) return a;
    assert (p >= 3);
    if (legendre_symbol(a, p) != +1) return -1;
    int b = 1;
    while (legendre_symbol(b, p) == 1) {
        b += 1;
    }
    int e = 0;
    int m = p - 1;
    while (m % 2 == 0) {
        m /= 2;
        e += 1;
    }
    ll x = powmod(a, (m - 1) / 2, p);
    ll y = a * x % p * x % p;
    x = x * a % p;
    ll z = powmod(b, m, p);
    while (y != 1) {
        int j = 0;
        for (ll t = y; t != 1; t = t * t % p) ++ j;
        if (e <= j) return -1;
        z = powmod(z, 1ll << (e - j - 1), p);
        x = x * z % p;
        z = z * z % p;
        y = y * z % p;
        e = j;
    }
    assert (x * x % p == a);
    return x;
}

vector<int> solve_modeqn(int a, int b, int c, int p) { // ax^2 + bx + c = 0 mod p
    int d = (b *(ll) b - 4ll * a * c) % p;
    if (d < 0) d += p;
    int w = modsqrt(d, p);
    if (w == -1) return vector<int>();
    vector<int> xs;
    xs.push_back((- b + w) *(ll) modinv(2 * a % p, p) % p);
    xs.push_back((- b - w) *(ll) modinv(2 * a % p, p) % p);
    if (xs[0] < 0) xs[0] += p;
    if (xs[1] < 0) xs[1] += p;
    whole(sort, xs);
    xs.erase(whole(unique, xs), xs.end());
    return xs;
}

int main() {
    vector<int> primes = list_primes(1e5);
    int p, r, q; scanf("%d%d%d", &p, &r, &q);
    repeat (i, q) {
        int a, b, c; scanf("%d%d%d", &a, &b, &c);
        vector<int> xs = solve_modeqn(a, b, c, p);
        if (xs.empty()) {
            printf("-1\n");
        } else {
            repeat (i, xs.size()) {
                printf("%d%c", xs[i], i + 1 == xs.size() ? '\n' : ' ');
            }
        }
    }
    return 0;
}