CTF感がすごい

## solution

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

## 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;
}