Yet Another Prefix Sum Problem

Let $X_i=A+iD$. If $X_i$ is divisible by $p$, then $X_i^{X_i}\equiv0\pmod p$. Otherwise, Fermat's little theorem allows us to reduce the exponent modulo $p-1$.

Group indices as $i=r+tp$, where $0\le r<\min(N,p)$. The number of indices in group $r$ is

$$m_r=\left\lfloor\frac{N-1-r}{p}\right\rfloor+1.$$

The integer $X_r$ may exceed the signed $64$-bit range, so compute only

$$b_r=(A\bmod p+(rD)\bmod p)\bmod p, \qquad e_r=(A\bmod(p-1)+(rD)\bmod(p-1))\bmod(p-1).$$

If $b_r=0$, every term in that group contributes zero. Otherwise, define

$$C_r=b_r^{e_r}\bmod p,\qquad R_r=b_r^{D\bmod(p-1)}\bmod p.$$

Then the term at index $i=r+tp$ satisfies

$$(A+iD)^{A+iD}\equiv C_rR_r^t\pmod p.$$

In particular, the first term is $b_r^{e_r}$, not necessarily $b_r^{b_r}$: congruence of the base modulo does not imply congruence of its exponent modulo .

The sum within a group is $C_rm_r$ if $R_r=1$, and otherwise

$$C_r\frac{R_r^{m_r}-1}{R_r-1}\pmod p.$$

Use a modular inverse for division. Since $b_r\ne0$, also $R_r\ne0$, so exponent reduction modulo is valid. Precompute inverses from to in linear time.

Add the contributions of every group. There are at most $\min(N,p)$ groups, and each requires $O(\log p)$ time for modular exponentiation. The overall time complexity is $O(p+\min(N,p)\log p)$ and the space complexity is . By reducing operands before multiplication, all intermediate values fit in signed -bit integers.

Solution written by GPT6