大好き！！ガウス整数の問題リストがまた 1 個増えた！

問題へのリンク

問題概要

1 つの正の整数 $R$ が与えられる。

二次元平面上の任意の格子点を頂点としたグラフにおいて、距離の平方が $R$ であるような 2 点間に辺を張っていく。

こうしてできたグラフの連結成分の個数を求めよ。無限個ある場合には "inf" と出力せよ。

(というテストケースが $T$ 個与えられる)

制約

$1 \le T \le 100$
$1 \le R \le 10^{9}$

最初の考察

一般にガウス整数 $z$ が存在して $z\bar{z} = R$ が成り立つとき、 $a, b$ を任意のガウス整数として

$az + bz$

で表されるような移動は可能となることがわかる。

たとえば $R = 5$ のときは、連結成分は 1 になる。なぜならば、ガウス整数の世界で 5 を素因数分解すると、

$R = 5 = (2 + i)(2 - i)$

となる。ここで、 $z = 2+i$ と $\bar{z} = 2-i$ は互いに素であるから、あるガウス整数 $a, b$ が存在して、

$a(2 + i) + b(2 - i) = 1$

が成立する。これはつまり、任意の格子点間の移動が可能であることを意味するのだ。

連結成分が 1 でない場合

では、連結成分が 1 でない場合はどんな場合だろうか。 $R = 2$ の場合はダメだ。 $R = z\bar{z}$ と表す方法は本質的には

$R = (1 + i)(1 - i)$

しかないが、 $1+i, 1-i$ は単数倍で移り合う関係なので、互いに素ではない。つまり、mod $1 + i$ の世界で等しいもの同士しか辺で結ばれないのだ。mod $1 + i$ では、すべてのガウス整数は 2 種類に類別されるので、答えは 2 ということになる。

他にも、 $R = 245 (= 5 \times 7^{2})$ という場合もある。この場合も $R = z\bar{z}$ と表す方法を考えたとき、 $z$ はからなず $7$ の倍数になる。つまり、mod $7$ の世界で等しいもの同士しか結ばれない。ガウス整数の世界では mod $7$ では、あまりが 49 個あるので、答えは 49 ということになる。

一般の場合

以上の考察を膨らませていくと、次の結論にいたる (ちゃんとした証明は tatyam さんの editorial より)。

$R$ を通常の整数の世界で素因数分解したとき

$4$ で割って $3$ 余る素因数であって、指数が奇数であるものが存在するとき："inf" (すべての格子点が孤立点になる)
そうでないとき：素因数のうち、 $2$ および $4$ で割って $3$ 余る部分のみを残して計算した値が答え

たとえば、 $R = 245 = 5 \times 7^{2}$ のときは $7^{2} = 49$ が答え。

計算量は $O(T\sqrt{R})$ となる。

コード

#include <bits/stdc++.h>
using namespace std;

vector<pair<long long, long long> > prime_factorize(long long n) {
    vector<pair<long long, long long> > res;
    for (long long p = 2; p * p <= n; ++p) {
        if (n % p != 0) continue;
        int num = 0;
        while (n % p == 0) { ++num; n /= p; }
        res.push_back(make_pair(p, num));
    }
    if (n != 1) res.push_back(make_pair(n, 1));
    return res;
}

void solve() {
    long long R;
    cin >> R;
    
    const auto &pf = prime_factorize(R);
    bool ok = true;
    long long res = 1;
    for (auto [p, e] : pf) {
        if (p % 4 == 3 && e % 2 == 1) ok = false;
        if (p % 4 == 1) continue;
        for (long long i = 0; i < e; ++i) res *= p;
    }
    
    if (ok) cout << res << endl;
    else cout << "inf" << endl;
}

int main() {
    int T;
    cin >> T;
    while (T--) solve();
}