1^k + 2^k + ... + n^k를 O(k)에 구하기

당연한 말이지만 $\mathbb{Z}_p$에서 생각한다.

Lagrange interpolation

구하려는 식은 $n$에 대한 $k+1$차 다항식이라는 사실이 알려져 있다. 따라서, $k+2$개의 함숫값만 알고 있으면 Lagrange interpolation을 통해 값을 계산할 수 있다. 함숫값을 구하는 건 거듭제곱을 해야 하기 때문에 $O(k \log k)$의 시간이 걸리고, $f(n)$을 계산하는 것은 역원을 나이브하게 계산할 경우 $O(k \log p)$의 시간이 걸린다. 사실 이 정도면 충분히 빠르다. TLE가 난다면 쓰레기 문제라는 뜻이기 때문에 풀지 않는 것을 권장한다.

 

Linear sieve

사실 $i^k$를 $i=0, 1, \ldots, k+1$에 대해 계산하는 것은 $O(k)$에 가능하다. $i^k$는 multiplicative function이라는 자명한 사실을 사용하면, $p^k$만 구해두면 충분하다는 것을 알 수 있다. $k$이하의 소수 개수는 $O\left( \frac{k}{\log k} \right)$이고, $p^k$를 구하는 데 걸리는 시간은 $O(\log k)$이므로 모든 $p^k$를 계산하는 것은 $O(k)$에 가능하다.

 

Interpolate_iota

Lagrange interpolation을 하기 위해 계산한 점의 x좌표는 $0$부터 $k+1$까지의 모든 정수다. 따라서, $f(n)$의 값을 계산하기 위해 필요한 값은 $i^{-1}$와 $(n-i)^{-1}$이다. 이를 빨리 계산하는 방법은 여러 가지가 있다. 내가 가장 간단하다고 생각하는 방법은 $(i!)^{-1}$을 모두 계산하는 방법이다. $(k!)^{-1}$을 계산하고 나면 $k-1, k-2, \ldots$을 하나씩 곱하면서 모든 팩토리얼의 역원을 구할 수 있게 된다. 여기에 $(i-1)!$을 곱하면 구하고자 하는 값을 얻게 된다. $(n-i)^{-1}$도 동일한 방식으로 $O(k + \log p)$에 전처리할 수 있다. $n-i$가 $0$이 되는 경우를 주의해서 코딩해야 한다.