abc 137 F - Polynomial Construction (python)

結局、フェルマーの小定理から x = i のとき 1 で、それ以外の時に 0 となるような p-1 次式が作れてしまうというだけの問題。賢い。
あとは a[ i ] = 1 のときに毎回 1 - (x-i)^(p-1) の係数を二項定理で作るだけ。

p = int(input())
a = list(map(int, input().split()))

MAX = 3000
fact = [1]*(MAX+1)
for i in range(1, MAX+1):
    fact[i] = (fact[i-1]*i) % p

inv = [1]*(MAX+1)
for i in range(2, MAX+1):
    inv[i] = inv[p % i]*(p-p//i) % p

fact_inv = [1]*(MAX+1)
for i in range(1, MAX+1):
    fact_inv[i] = fact_inv[i-1] * inv[i] % p

def comb(n, k):
    if n < k:
        return 0
    return fact[n] * fact_inv[n-k] * fact_inv[k] % p

ans = [0]*p
for i,ai in enumerate(a):
  if ai == 1:
    ans[-1] += 1
    for j in range(p):
      ans[j] -= pow(-i,j,p)*comb(p-1,j)
      ans[j] %= p
print(*ans[::-1])