Yet Another Prefix Sum Problem

$X_i=A+iD$ とおく。$X_i$ が $p$ の倍数なら $X_i^{X_i}\equiv0\pmod p$ である。それ以外では、フェルマーの小定理により指数を $p-1$ で割った余りに置き換えられる。

添字を $i=r+tp$ によって分類する。$0\le r<\min(N,p)$ であり、グループ $r$ の項数は

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

である。

$X_r$ は符号付き $64$ ビット整数の範囲を超えるため、次の剰余のみを計算する。

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

$b_r=0$ なら、そのグループの寄与はすべて $0$ である。それ以外では

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

とおく。このとき $i=r+tp$ の項は

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

となる。最初の項は一般に $b_r^{b_r}$ ではなく $b_r^{e_r}$ である。底の に関する剰余が同じでも、指数の に関する剰余は一致しないためである。

グループの和は $R_r=1$ なら $C_rm_r$、それ以外では

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

である。除算には逆元を用いる。$b_r\ne0$ なので $R_r\ne0$ であり、指数は で割った余りに縮められる。 から の逆元は線形時間で前計算できる。

グループの和を合計する。グループ数は高々 $\min(N,p)$ で、各グループは $O(\log p)$ 時間で計算できる。合計の時間計算量は $O(p+\min(N,p)\log p)$、空間計算量は である。積の計算前に剰余を取れば、すべての中間値は符号付き ビット整数に収まる。

Solution written by GPT6