Problem p02539 (python)

Fixed submission:
import sys
input = sys.stdin.readline
sys.setrecursionlimit(1000000)

from collections import defaultdict
from collections import deque
import heapq

MOD = 998244353

def DD(arg): return defaultdict(arg)

def inv(n): return pow(n, MOD-2, MOD)

kaijo_memo = []
def kaijo(n):
  if(len(kaijo_memo) > n): return kaijo_memo[n]
  if(len(kaijo_memo) == 0): kaijo_memo.append(1)
  while(len(kaijo_memo) <= n): kaijo_memo.append(kaijo_memo[-1] * len(kaijo_memo) % MOD)
  return kaijo_memo[n]

gyaku_kaijo_memo = []
def gyaku_kaijo(n):
  if(len(gyaku_kaijo_memo) > n): return gyaku_kaijo_memo[n]
  if(len(gyaku_kaijo_memo) == 0): gyaku_kaijo_memo.append(1)
  while(len(gyaku_kaijo_memo) <= n): gyaku_kaijo_memo.append(gyaku_kaijo_memo[-1] * pow(len(gyaku_kaijo_memo),MOD-2,MOD) % MOD)
  return gyaku_kaijo_memo[n]

ROOT = 3
MOD = 998244353
roots  = [pow(ROOT,(MOD-1)>>i,MOD) for i in range(24)] # 1 の 2^i 乗根
iroots = [pow(x,MOD-2,MOD) for x in roots] # 1 の 2^i 乗根の逆元
 
def untt(a,n):
    for i in range(n):
        m = 1<<(n-i-1)
        for s in range(1<<i):
            w_N = 1
            s *= m*2
            for p in range(m):
                a[s+p], a[s+p+m] = (a[s+p]+a[s+p+m])%MOD, (a[s+p]-a[s+p+m])*w_N%MOD
                w_N = w_N*roots[n-i]%MOD

def iuntt(a,n):
    for i in range(n):
        m = 1<<i
        for s in range(1<<(n-i-1)):
            w_N = 1
            s *= m*2
            for p in range(m):
                a[s+p], a[s+p+m] = (a[s+p]+a[s+p+m]*w_N)%MOD, (a[s+p]-a[s+p+m]*w_N)%MOD
                w_N = w_N*iroots[i+1]%MOD
            
    inv = pow((MOD+1)//2,n,MOD)
    for i in range(1<<n):
        a[i] = a[i]*inv%MOD

def convolution(a,b):
    la = len(a)
    lb = len(b)
    deg = la+lb-2
    n = deg.bit_length()
    if min(la, lb) <= 50:
        if la < lb:
            la,lb = lb,la
            a,b = b,a
        res = [0]*(la+lb-1)
        for i in range(la):
            for j in range(lb):
                res[i+j] += a[i]*b[j]
                res[i+j] %= MOD
        return res
 
    N = 1<<n
    a += [0]*(N-len(a))
    b += [0]*(N-len(b))
    untt(a,n)
    untt(b,n)
    for i in range(N):
      a[i] = a[i]*b[i]%MOD
    iuntt(a,n)
    return a[:deg+1]

N = int(input())

nokori = [1]
for n in range(1,2*N+10):
  nokori.append((nokori[-1]*(2*n-1)*(2*n)*inv(2)*inv(n))%MOD)

dic = DD(int)
for _ in range(N*2):
  dic[int(input())] += 1
A = [dic[k] for k in dic]
  
P = []
for a in A:
  temp = 1
  count = 0
  P.append([temp])
  while a >= 2:
    temp *= a*(a-1)//2
    temp %= MOD
    count += 1
    P[-1].append((temp*gyaku_kaijo(count))%MOD)
    a -= 2

Q = []
for i in range(len(P)):
  heapq.heappush(Q,(len(P[i]),i))

while len(Q)>1:
  _,p = heapq.heappop(Q)
  _,q = heapq.heappop(Q)
  temp = convolution(P[p], P[q])
  P.append(temp)
  heapq.heappush(Q,(len(temp),len(P)-1))
Q = list(P[-1])

ans = 0
M = len(Q)

def sign(x):
  if x%2: return -1
  else: return 1

for i in range(M):
  ans += sign(i) * Q[i] * nokori[N-i]
  ans %= MOD
  
print(ans)

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline
sys.setrecursionlimit(1000000)

from collections import defaultdict
from collections import deque
import heapq
from scipy import fft
import numba
from numba import njit, b1, i4, i8, f8

MOD = 998244353

def DD(arg): return defaultdict(arg)

def inv(n): return pow(n, MOD-2, MOD)

kaijo_memo = []
def kaijo(n):
  if(len(kaijo_memo) > n): return kaijo_memo[n]
  if(len(kaijo_memo) == 0): kaijo_memo.append(1)
  while(len(kaijo_memo) <= n): kaijo_memo.append(kaijo_memo[-1] * len(kaijo_memo) % MOD)
  return kaijo_memo[n]

gyaku_kaijo_memo = []
def gyaku_kaijo(n):
  if(len(gyaku_kaijo_memo) > n): return gyaku_kaijo_memo[n]
  if(len(gyaku_kaijo_memo) == 0): gyaku_kaijo_memo.append(1)
  while(len(gyaku_kaijo_memo) <= n): gyaku_kaijo_memo.append(gyaku_kaijo_memo[-1] * pow(len(gyaku_kaijo_memo),MOD-2,MOD) % MOD)
  return gyaku_kaijo_memo[n]

import numpy as np

@njit((i8, i8), cache=True)
def ntt_precompute(N, primitive_root):
    roots = np.empty(1 << N, np.int64)
    iroots = np.empty(1 << N, np.int64)
    roots[0] = iroots[0] = 1
    for n in range(N):
        x = mpow(primitive_root, (MOD - 1) >> n + 2)
        y = mpow(x, MOD - 2)
        for i in range(1 << n):
            roots[(1 << n) | i] = roots[i] * x % MOD
            iroots[(1 << n) | i] = iroots[i] * y % MOD
    return roots, iroots

@njit((i8, i8), cache=True)
def mpow(a, n):
    p = 1
    while n:
        if n & 1:
            p = p * a % MOD
        a = a * a % MOD
        n >>= 1
    return p

@njit((i8[:], i8[:], i8[:], numba.b1), cache=True)
def ntt(roots, iroots, A, inverse):
    N = len(A)
    if not inverse:
        m = N >> 1
        while m:
            for k in range(N // m >> 1):
                s = 2 * m * k
                for i in range(s, s + m):
                    x, y = A[i], A[i + m] * roots[k] % MOD
                    A[i], A[i + m] = x + y, x - y
            A %= MOD
            m >>= 1
    else:
        m = 1
        while m < N:
            for k in range(N // m >> 1):
                s = 2 * m * k
                for i in range(s, s + m):
                    x, y = A[i], A[i + m]
                    A[i], A[i + m] = x + y, (x - y) * iroots[k] % MOD
            A %= MOD
            m <<= 1
        invN = mpow(N, MOD - 2)
        for i in range(N):
            A[i] = A[i] * invN % MOD
    return A
 
 
@njit((i8[:], i8[:], i8[:], i8[:]), cache=True)
def ntt_convolve(roots, iroots, A, B):
    n = 0
    L = len(A) + len(B) - 1
    while (1 << n) < L:
        n += 1
    A1 = np.zeros(1 << n, np.int64)
    A1[:len(A)] = A
    B1 = np.zeros(1 << n, np.int64)
    B1[:len(B)] = B
    A, B = A1, B1
    ntt(roots, iroots, A, False)
    ntt(roots, iroots, B, False)
    for i in range(len(A)):
        A[i] = A[i] * B[i] % MOD
    ntt(roots, iroots, A, True)
    return A[:L]
  
N = int(input())

max_ntt = 18
roots, iroots = ntt_precompute(max_ntt, 3)

nokori = [1]
for n in range(1,2*N+10):
  nokori.append((nokori[-1]*(2*n-1)*(2*n)*inv(2)*inv(n))%MOD)

dic = DD(int)
for _ in range(N*2):
  dic[int(input())] += 1
A = [dic[k] for k in dic]
  
P = []
for a in A:
  temp = 1
  count = 0
  P.append([temp])
  while a >= 2:
    temp *= a*(a-1)//2
    temp %= MOD
    count += 1
    P[-1].append((temp*gyaku_kaijo(count))%MOD)
    a -= 2


data = [np.array(p) for p in P]
Q = []
for i in range(len(data)):
  heapq.heappush(Q,(len(data[i]),i))

while len(Q)>1:
  _,p = heapq.heappop(Q)
  _,q = heapq.heappop(Q)
  temp = ntt_convolve(roots, iroots, data[p], data[q])
  data.append(temp)
  heapq.heappush(Q,(len(temp),len(data)-1))
Q = list(data[-1])

ans = 0
M = len(Q)

def sign(x):
  if x%2: return -1
  else: return 1

for i in range(M):
  ans += sign(i) * Q[i] * nokori[N-i]
  ans %= MOD
  
print(ans)

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline
sys.setrecursionlimit(1000000)

from collections import defaultdict
from collections import deque
import heapq
from scipy import fft

MOD = 998244353

def DD(arg): return defaultdict(arg)

def inv(n): return pow(n, MOD-2, MOD)

kaijo_memo = []
def kaijo(n):
  if(len(kaijo_memo) > n): return kaijo_memo[n]
  if(len(kaijo_memo) == 0): kaijo_memo.append(1)
  while(len(kaijo_memo) <= n): kaijo_memo.append(kaijo_memo[-1] * len(kaijo_memo) % MOD)
  return kaijo_memo[n]

gyaku_kaijo_memo = []
def gyaku_kaijo(n):
  if(len(gyaku_kaijo_memo) > n): return gyaku_kaijo_memo[n]
  if(len(gyaku_kaijo_memo) == 0): gyaku_kaijo_memo.append(1)
  while(len(gyaku_kaijo_memo) <= n): gyaku_kaijo_memo.append(gyaku_kaijo_memo[-1] * pow(len(gyaku_kaijo_memo),MOD-2,MOD) % MOD)
  return gyaku_kaijo_memo[n]

import numpy as np

@njit((i8, i8), cache=True)
def ntt_precompute(N, primitive_root):
    roots = np.empty(1 << N, np.int64)
    iroots = np.empty(1 << N, np.int64)
    roots[0] = iroots[0] = 1
    for n in range(N):
        x = mpow(primitive_root, (MOD - 1) >> n + 2)
        y = mpow(x, MOD - 2)
        for i in range(1 << n):
            roots[(1 << n) | i] = roots[i] * x % MOD
            iroots[(1 << n) | i] = iroots[i] * y % MOD
    return roots, iroots

@njit((i8, i8), cache=True)
def mpow(a, n):
    p = 1
    while n:
        if n & 1:
            p = p * a % MOD
        a = a * a % MOD
        n >>= 1
    return p

@njit((i8[:], i8[:], i8[:], numba.b1), cache=True)
def ntt(roots, iroots, A, inverse):
    N = len(A)
    if not inverse:
        m = N >> 1
        while m:
            for k in range(N // m >> 1):
                s = 2 * m * k
                for i in range(s, s + m):
                    x, y = A[i], A[i + m] * roots[k] % MOD
                    A[i], A[i + m] = x + y, x - y
            A %= MOD
            m >>= 1
    else:
        m = 1
        while m < N:
            for k in range(N // m >> 1):
                s = 2 * m * k
                for i in range(s, s + m):
                    x, y = A[i], A[i + m]
                    A[i], A[i + m] = x + y, (x - y) * iroots[k] % MOD
            A %= MOD
            m <<= 1
        invN = mpow(N, MOD - 2)
        for i in range(N):
            A[i] = A[i] * invN % MOD
    return A
 
 
@njit((i8[:], i8[:], i8[:], i8[:]), cache=True)
def ntt_convolve(roots, iroots, A, B):
    n = 0
    L = len(A) + len(B) - 1
    while (1 << n) < L:
        n += 1
    A1 = np.zeros(1 << n, np.int64)
    A1[:len(A)] = A
    B1 = np.zeros(1 << n, np.int64)
    B1[:len(B)] = B
    A, B = A1, B1
    ntt(roots, iroots, A, False)
    ntt(roots, iroots, B, False)
    for i in range(len(A)):
        A[i] = A[i] * B[i] % MOD
    ntt(roots, iroots, A, True)
    return A[:L]
  
N = int(input())

max_ntt = 18
roots, iroots = ntt_precompute(max_ntt, 3)

nokori = [1]
for n in range(1,2*N+10):
  nokori.append((nokori[-1]*(2*n-1)*(2*n)*inv(2)*inv(n))%MOD)

dic = DD(int)
for _ in range(N*2):
  dic[int(input())] += 1
A = [dic[k] for k in dic]
  
P = []
for a in A:
  temp = 1
  count = 0
  P.append([temp])
  while a >= 2:
    temp *= a*(a-1)//2
    temp %= MOD
    count += 1
    P[-1].append((temp*gyaku_kaijo(count))%MOD)
    a -= 2


data = [np.array(p) for p in P]
Q = []
for i in range(len(data)):
  heapq.heappush(Q,(len(data[i]),i))

while len(Q)>1:
  _,p = heapq.heappop(Q)
  _,q = heapq.heappop(Q)
  temp = ntt_convolve(roots, iroots, data[p], data[q])
  data.append(temp)
  heapq.heappush(Q,(len(temp),len(data)-1))
Q = list(data[-1])

ans = 0
M = len(Q)

def sign(x):
  if x%2: return -1
  else: return 1

for i in range(M):
  ans += sign(i) * Q[i] * nokori[N-i]
  ans %= MOD
  
print(ans)

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline
sys.setrecursionlimit(1000000)

from collections import defaultdict
from collections import deque
import heapq
from scipy import fft

MOD = 998244353

def DD(arg): return defaultdict(arg)

def inv(n): return pow(n, MOD-2, MOD)

kaijo_memo = []
def kaijo(n):
  if(len(kaijo_memo) > n): return kaijo_memo[n]
  if(len(kaijo_memo) == 0): kaijo_memo.append(1)
  while(len(kaijo_memo) <= n): kaijo_memo.append(kaijo_memo[-1] * len(kaijo_memo) % MOD)
  return kaijo_memo[n]

gyaku_kaijo_memo = []
def gyaku_kaijo(n):
  if(len(gyaku_kaijo_memo) > n): return gyaku_kaijo_memo[n]
  if(len(gyaku_kaijo_memo) == 0): gyaku_kaijo_memo.append(1)
  while(len(gyaku_kaijo_memo) <= n): gyaku_kaijo_memo.append(gyaku_kaijo_memo[-1] * pow(len(gyaku_kaijo_memo),MOD-2,MOD) % MOD)
  return gyaku_kaijo_memo[n]

import numpy as np

@njit((i8, i8), cache=True)
def ntt_precompute(N, primitive_root):
    roots = np.empty(1 << N, np.int64)
    iroots = np.empty(1 << N, np.int64)
    roots[0] = iroots[0] = 1
    for n in range(N):
        x = mpow(primitive_root, (MOD - 1) >> n + 2)
        y = mpow(x, MOD - 2)
        for i in range(1 << n):
            roots[(1 << n) | i] = roots[i] * x % MOD
            iroots[(1 << n) | i] = iroots[i] * y % MOD
    return roots, iroots

@njit((i8, i8), cache=True)
def mpow(a, n):
    p = 1
    while n:
        if n & 1:
            p = p * a % MOD
        a = a * a % MOD
        n >>= 1
    return p

@njit((i8[:], i8[:], i8[:], numba.b1), cache=True)
def ntt(roots, iroots, A, inverse):
    N = len(A)
    if not inverse:
        m = N >> 1
        while m:
            for k in range(N // m >> 1):
                s = 2 * m * k
                for i in range(s, s + m):
                    x, y = A[i], A[i + m] * roots[k] % MOD
                    A[i], A[i + m] = x + y, x - y
            A %= MOD
            m >>= 1
    else:
        m = 1
        while m < N:
            for k in range(N // m >> 1):
                s = 2 * m * k
                for i in range(s, s + m):
                    x, y = A[i], A[i + m]
                    A[i], A[i + m] = x + y, (x - y) * iroots[k] % MOD
            A %= MOD
            m <<= 1
        invN = mpow(N, MOD - 2)
        for i in range(N):
            A[i] = A[i] * invN % MOD
    return A
 
 
@njit((i8[:], i8[:], i8[:], i8[:]), cache=True)
def ntt_convolve(roots, iroots, A, B):
    n = 0
    L = len(A) + len(B) - 1
    while (1 << n) < L:
        n += 1
    A1 = np.zeros(1 << n, np.int64)
    A1[:len(A)] = A
    B1 = np.zeros(1 << n, np.int64)
    B1[:len(B)] = B
    A, B = A1, B1
    ntt(roots, iroots, A, False)
    ntt(roots, iroots, B, False)
    for i in range(len(A)):
        A[i] = A[i] * B[i] % MOD
    ntt(roots, iroots, A, True)
    return A[:L]
  
N = int(input())

max_ntt = 18
roots, iroots = ntt_precompute(max_ntt, 3)

nokori = [1]
for n in range(1,2*N+10):
  nokori.append((nokori[-1]*(2*n-1)*(2*n)*inv(2)*inv(n))%MOD)

dic = DD(int)
for _ in range(N*2):
  dic[int(input())] += 1
A = [dic[k] for k in dic]
  
P = []
for a in A:
  temp = 1
  count = 0
  P.append([temp])
  while a >= 2:
    temp *= a*(a-1)//2
    temp %= MOD
    count += 1
    P[-1].append((temp*gyaku_kaijo(count))%MOD)
    a -= 2


data = [np.array(p) for p in P]
Q = []
for i in range(len(data)):
  heapq.heappush(Q,(len(data[i]),i))

while len(Q)>1:
  _,p = heapq.heappop(Q)
  _,q = heapq.heappop(Q)
  temp = ntt_convolve(roots, iroots, data[p], data[q])
  data.append(temp)
  heapq.heappush(Q,(len(temp),len(data)-1))
Q = list(data[-1])

ans = 0
M = len(Q)

def sign(x):
  if x%2: return -1
  else: return 1

for i in range(M):
  ans += sign(i) * Q[i] * nokori[N-i]
  ans %= MOD
  
print(ans)