Problem p02666 (python)

Fixed submission:
import sys;input=sys.stdin.readline
class UnionFindNode:
    def __init__(self, group_id, parent=None, value=None):
        self.group_id_ = group_id
        self.parent_ = parent
        self.value = value
        self.rank_ = 1
        self.member_num_ = 1

    def is_root(self):
        return not self.parent_

    def root(self):
        parent = self
        while not parent.is_root():
            parent = parent.parent_
            self.parent_ = parent
        return parent

    def find(self):
        root = self.root()
        return root.group_id_

    def rank(self):
        root = self.root()
        return root.rank_

    def unite(self, unite_node):
        root = self.root()
        unite_root = unite_node.root()

        if root.group_id_ != unite_root.group_id_:
            if root.rank() > unite_root.rank():
                n_root, child = root, unite_root
            else:
                n_root, child = unite_root, root
            child.parent_ = n_root
            n_root.rank_ = max(n_root.rank_, child.rank_ + 1)
            n_root.member_num_ = n_root.member_num_ + child.member_num_

mod = 10**9+7
NNN = (10**6)
g1 = [1, 1]
for i in range( 2, NNN + 1 ):
    g1.append( ( g1[-1] * i ) % mod )
N, = map(int, input().split())
P = list(map(int, input().split()))
K = P.count(-1)

nodes = [UnionFindNode(i) for i in range(N)]
for i, p in enumerate(P):
    p-=1
    if p >= 0:
        nodes[p].unite(nodes[i])

rs = []
X = [0] * N
for i in range(N):
    if nodes[i].is_root():
        rs.append(i)
    if P[i] == -1:
        X[nodes[i].find()] = 1

M = len(rs)
R = N*pow(N-1, K, mod)%mod

Y = []
for i in range(M):
    if X[rs[i]]:
        Y.append(nodes[rs[i]].member_num_)
        R = (R-(nodes[rs[i]].member_num_-1)*pow(N-1, K-1, mod))%mod
    else:
        R = (R-pow(N-1, K, mod))%mod
nY = len(Y)
dp = [[0]*(nY+1) for _ in range(nY+1)]
dp[0][0] = 1
for i in range(nY):
    for j in range(nY):
        dp[i+1][j] = (dp[i][j] * Y[i] + (dp[i][j-1] if j > 0 else 0)) % mod
for i in range(nY):
    if nY-i == 1:
        continue
#    print(nY-i, dp[-1][i])
    R -= pow(N-1, K-(nY-i), mod) * g1[nY-i-1] * dp[-1][i]
    R %= mod
print(R)

Buggy submission (verdict: Runtime Error):
import sys
import numpy as np

read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines

MOD = 998244353
def fft_convolve(f, g, MOD=MOD):
    fft = np.fft.rfft
    ifft = np.fft.irfft
    Lf = f.shape[-1]
    Lg = g.shape[-1]
    L = Lf + Lg - 1
    fft_len = 1 << L.bit_length()
    fl = f & (1 << 15) - 1
    fh = f >> 15
    gl = g & (1 << 15) - 1
    gh = g >> 15

    def conv(f, g):
        return ifft(fft(f, fft_len) * fft(g, fft_len))[:, :L]

    x = conv(fl, gl) % MOD
    y = conv(fl + fh, gl + gh) % MOD
    z = conv(fh, gh) % MOD
    a, b, c = map(lambda x: (x + .5).astype(np.int64), [x, y, z])
    return (a + ((b - a - c) << 15) + (c << 30)) % MOD

class UnionFindNode:
    def __init__(self, group_id, parent=None, value=None):
        self.group_id_ = group_id
        self.parent_ = parent
        self.value = value
        self.rank_ = 1
        self.member_num_ = 1

    def is_root(self):
        return not self.parent_

    def root(self):
        parent = self
        while not parent.is_root():
            parent = parent.parent_
            self.parent_ = parent
        return parent

    def find(self):
        root = self.root()
        return root.group_id_

    def rank(self):
        root = self.root()
        return root.rank_

    def unite(self, unite_node):
        root = self.root()
        unite_root = unite_node.root()

        if root.group_id_ != unite_root.group_id_:
            if root.rank() > unite_root.rank():
                n_root, child = root, unite_root
            else:
                n_root, child = unite_root, root
            child.parent_ = n_root
            n_root.rank_ = max(n_root.rank_, child.rank_ + 1)
            n_root.member_num_ = n_root.member_num_ + child.member_num_

mod = 10**9+7
NNN = (10**6)
g1 = [1, 1]
for i in range( 2, NNN + 1 ):
    g1.append( ( g1[-1] * i ) % mod )
N, = map(int, input().split())
P = list(map(int, input().split()))
K = P.count(-1)

nodes = [UnionFindNode(i) for i in range(N)]
for i, p in enumerate(P):
    p-=1
    if p >= 0:
        nodes[p].unite(nodes[i])

rs = []
X = [0] * N
for i in range(N):
    if nodes[i].is_root():
        rs.append(i)
    if P[i] == -1:
        X[nodes[i].find()] = 1

M = len(rs)
R = N*pow(N-1, K, mod)%mod

Y = []
for i in range(M):
    if X[rs[i]]:
        Y.append(nodes[rs[i]].member_num_)
        R = (R-(nodes[rs[i]].member_num_-1)*pow(N-1, K-1, mod))%mod
    else:
        R = (R-pow(N-1, K, mod))%mod
nY = len(Y)
A = np.array(Y, np.int64)
polys = np.empty((nY, 2), np.int64)
polys[:, 1] = 1
polys[:, 0] = A

while len(polys) > 1:
    if len(polys) & 1:
        polys = np.pad(polys, ((0, 1), (0, 0)))
        polys[-1, 0] = 1
    P = polys[:len(polys) // 2]
    Q = polys[len(polys) // 2:]
    polys = fft_convolve(P, Q)

P = polys[0]
for i in range(nY):
    if nY-i == 1:
        continue
    R -= pow(N-1, K-(nY-i), mod) * g1[nY-i-1] * P[i]
    R %= mod
print(R)