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

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

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

これ 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) の組の個数を求めよ。

制約

考えたこと

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

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