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)$이고, 해당 그룹의 항 수는

$$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