これ FFT を使えば $P \le 10^{5}$ でもできるね。実際は $O(P^{2})$ で間に合う。

問題概要

正の素数 $P$ と、正の整数 $N$ が与えられる。

$x^{N} + y^{N} \equiv z^{N} \pmod P$
$x, y, z$ は $0$ 以上 $P-1$ 以下の整数

を満たすような $(x, y, z)$ の組の個数を求めよ。

制約

$2 \le P \le 10^{4}$
$P$ は素数
$1 \le N \le 10^{4}$

考えたこと

まずは次のように集計処理をしよう。このように、「〜が何個あるか」という情報に変換して考えるのは典型！！

num[v] ← $x^{N} \equiv v \pmod P$ となるような整数 $x$ ( $0 \le x \le P-1$ ) の個数

これは繰り返し二乗法で計算すれば $O(P \log N)$ でできる。

統合する

配列 num が求められたら、それを使って次の配列を求めよう。

num2[v] ← $x^{N} + y^{N} \equiv v \pmod P$ となるような整数 $(x, y)$ の組の個数

これは単純には次のような処理で $O(P^{2})$ で求められる！

vector<long long> num2(P, 0);
for (int u = 0; u < P; ++u) {
    for (int v = 0; v < P; ++v) {
        num2[(u+v)%P] += num[u] + num[v];
    }
}

これを求めておけば、あとは答えは次のように求められる。

long long res = 0;
for (int v = 0; v < P; ++v) {
    res += num2[v] + num[v];
}

コード

計算量は $O(P \log N + P^{2})$ となる。

なお、num2 を求める部分は、FFT を用いれば $O(P \log P)$ でできる。よってこの問題は $O(P (\log P + \log N) )$ でもできるので、 $P \le 10^{5}$ , $N \le 10^{18}$ でも解ける。

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

long long modpow(long long a, long long n, long long mod) {
    long long res = 1;
    while (n > 0) {
        if (n & 1) res = res * a % mod;
        a = a * a % mod;
        n >>= 1;
    }
    return res;
}

int main() {
    long long P, N;
    cin >> P >> N;
    vector<long long> num(P, 0), num2(P, 0);
    for (int v = 0; v < P; ++v) num[modpow(v, N, P)]++;
    for (int u = 0; u < P; ++u) {
        for (int v = 0; v < P; ++v) {
            int s = (u + v >= P ? u + v - P : u + v);
            num2[s] += num[u] * num[v];
        }
    }
    long long res = 0;
    for (int v = 0; v < P; ++v) res += num2[v] * num[v];
    cout << res << endl;
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

JOI 春合宿 2007 day2-2 Fermat (難易度 7)

問題概要

制約

考えたこと

統合する

コード