n차 다항식 함숫값 빠르게 계산하기(페르마의 소정리로 거듭제곱 빠르게 계산하기)

1. 문제

 

25025번: 다항식 계산 (acmicpc.net)

25025번: 다항식 계산

 

2. 풀이

 

결국에 구해야하는 값은 $a_{i}k^{i} mod P$이고 k와 P가 서로소이므로 페르마의 소정리를 이용할 수 있다.

 

그런데 1초안에 P개의 값을 계산해야하니 이거 쉽지 않다

 

계수는 최악의 경우 $10^{6}$개이고.. 매 함숫값마다 이들을 모두 순회해야하니.. $O(10^{9})$는 돌아야겠는데 

 

이거 1초안에 당연히 안될거고..

 

근데 조금 생각해보면...

 

페르마의 소정리에 의해 $$k^{p-1} ≡ 1 (mod p)$$

 

하지만 N이 p보다 클수 있는데 그런 경우에 $k^{N} mod p$는 어떻게 구해야하는가?

 

합동식은 정수 q번 거듭제곱해도 성립하므로...

 

p-1에서 q번 거듭제곱한다면.. $(k^{p-1})^{q} ≡ 1 (mod p)$

 

여기에 $k^{r}$을 양변에 곱해도 합동식이 성립하므로... $$k^{(p-1)q+r} ≡ k^{r} (mod p)$$

 

만약에, $N = (p-1)q + r$이 성립한다고 한다면... 이 식이 무엇을 뜻하는가?

 

N을 p-1로 나눈 몫이 q이고 나머지가 r이다.

 

즉, $k^{N} mod p$는 페르마의 소정리를 이용한다면.. N을 p-1로 나눈 나머지 r에 대하여 $k^{r} mod p$로 압축할수 있다

 

그래서 먼저 $O(N)$에 1부터 N까지 순회하여 p-1로 나눈 나머지를 모두 구해준다.

 

그러면 1부터 N까지가 0 ~ p-2중 하나로 압축될 것이며, 그러면서 동일한 거듭제곱 항은 서로 묶어서 계산할 수 있다

 

그래서 동일한 거듭제곱 항끼리 계수를 서로 묶어서 미리 계산해준다.

 

이렇게 한다면, 모든 i = 0,1,2,...,p-1에 대하여 f(i)를 계산할때, 계산해야할 거듭제곱 항은 0 ~ p-2까지로 압축되며,

 

$O(N+p^{2})$에 문제를 해결할 수 있다

 

약 $2*10^{6}$이라 1초안에 돌수 있을듯

 

from sys import stdin

n,p = map(int,stdin.readline().split())

A = list(map(int,stdin.readline().split()))

A_dict = {}

f1 = 0

#페르마의 소정리에 의해, k^n mod p = k^(n%p) mod p
#차수가 같은 항끼리 묶어서 계수를 미리 계산해 N+1 > p-1개의 항으로 압축
for i in range(n+1):
    
    A[i] %= p

    f1 += A[i]

    r = (n-i) % (p-1)

    A_dict[r] = A_dict.get(r,0) + A[i]
    A_dict[r] %= p

print(A[-1]) #f(0) mod p
print(f1 % p) #f(1) mod p

#i = 2,3,...,p-1에 대해 f(i) mod p를 계산
#A_dict에 차수가 k=0,1,2,..p-1인 항의 계수 v가 계산되어있다
#k 차수 항은 v*i^k mod p로 계산 가능
#분할정복을 이용한 거듭제곱으로 빠르게
for i in range(2,p):

    answer = 0

    for k,v in A_dict.items():
        
        answer += pow(i,k,p)*v
        answer %= p
    
    print(answer)