はじめに、前準備として $0^n,1^n,2^n,\cdots,(p-1)^n$ を $\bmod p$ で求めることを考えます。これは繰り返し二乗法を用いることで一つにつき $O(\log n)$ で求めることができるので、この処理は $O(p \log n)$ で可能です。

また、この結果を用いて $a^n \equiv i$ かつ $0 \leq a \leq p-1$ となる $a$ の個数 ( $cnt[i]$ とします) を $0 \leq i \leq p-1$ について求めることが $O(p)$ で可能です。

\(O(p^2+p \log n)\) 解法

$x,y$ を決め打った際、 $z$ として適するのは $cnt[(x^n+y^n) \bmod p]$ 個あります。また、 $x,y$ としてありえるのは $p^2$ 個なので、 $x,y$ を全探索することで前準備と合わせて $O(p^2+p \log n)$ で解くことができました。

現在ではこの計算量でも AC することができますが、出題当時では環境による違いで不可能だったようです。本解説ではより高速な解法を紹介します。

\(O(p \log p+p \log n)\) 解法

$z$ を決め打つことを考えます。この時、 $x,y$ として適する組の数は $cnt[0] \times cnt[z^n \bmod p] + cnt[1] \times cnt[(z^n-1) \bmod p] + \cdots + cnt[z^n \bmod p] \times cnt[0]$ 個です。

また、畳み込みにより $cnt2[i] = cnt[0] \times cnt[i] + cnt[1] \times cnt[i-1] + \cdots + cnt[i] \times cnt[0]$ となる配列 $cnt2$ を $O(p \log p)$ で入手できます。

よって、この問題は $O(p \log p + p \log n)$ で解くことができました。

\(O(p \log n)\) 解法

$x^n+y^n \equiv z^n \equiv i$ となるような $(x,y,z)$ の数を $sum[i]$ とします。この時、 $i>1$ に対して $sum[i]=sum[1]$ が成り立つことを証明します。ただし、 $cnt[i]>0$ とし、 $a \equiv x^n,b \equiv y^n,c \equiv z^n$ とします。

$c \equiv 1$ の時、 $x^n+y^n \equiv 1$ です。 $c>1$ となる $z$ を $t$ として、両辺に $t^n$ をかけると $(xt)^n+(yt)^n \equiv t^n$ がいえます。よって $(xt,yt,t)$ は条件を満たす組であり、 $(x,y) \neq (x',y')$ ならば $(xt,yt) \neq (x't,y't)$ であるため $sum[1] \leq sum[i] (i>1)$ が示せます。

また、 $d \equiv z^{-1}$ として $x^n+y^n \equiv z^n$ の場合、両辺に $d^n$ を かけることで $(xd)^n+(yd)^n \equiv 1$ なので、 $(xd,yd,1)$ は必ず条件を満たします。 $(x,y) \neq (x',y')$ ならば $(xd,yd) \neq (x'd,y'd)$ であるため $sum[i] \leq sum[1] (i>1)$ が示せます。

以上より \(sum[i]=sum[1] (i>1)\) が示せました。よって、 \(sum[0],sum[1]\) と \(z^n\) が取りうる値の種類をそれぞれ \(O(p)\) で求めることで、この問題は \(O(p \log n)\) で解くことができました。