Problem p02965 (python)

Fixed submission:
import sys
input = sys.stdin.readline

import numpy as np

N,M = map(int,input().split())
MOD = 998244353

def cumprod_mod(arr):
    L = len(arr)
    Lsq = int(L**.5+1)
    arr = np.resize(arr,Lsq**2).reshape(Lsq,Lsq)
    for n in range(1,Lsq):
        arr[:,n] *= arr[:,n-1]; arr[:,n] %= MOD
    for n in range(1,Lsq):
        arr[n] *= arr[n-1,-1]; arr[n] %= MOD
    arr = arr.ravel()[:L]
    return arr

U = 5*10**6
x = np.arange(U,dtype=np.int64)
x[0] = 1
fact = cumprod_mod(x)

x = np.arange(U,0,-1,dtype=np.int64)
x[0] = pow(int(fact[-1]),MOD-2,MOD)
fact_inv = cumprod_mod(x)[::-1]

def comb(n,k):
    return fact[n] * fact_inv[k] % MOD * fact_inv[n-k] % MOD

x = comb(3*M+N-1,N-1) # 何でもいいから作る
y = N * comb(M+N-2,N-1) # どれかが 2M+1 以上

# 奇数がM+1個以上あるものを除く
odd_cnt = np.arange(M+2,3*M+1,2)
S = (3*M - odd_cnt)//2
bl = (odd_cnt <= N)&(S>=0)
odd_cnt = odd_cnt[bl]
S = S[bl]
arr = fact[N] * fact_inv[odd_cnt] % MOD * fact_inv[N-odd_cnt] % MOD
# SをN個で作る -> S+N-1 C N-1
arr *= fact[S+N-1] * fact_inv[N-1] % MOD * fact_inv[S] % MOD
arr %= MOD
z = arr.sum() % MOD

answer = x-y-z
answer %= MOD
print(answer)

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

import numpy as np

N,M = map(int,input().split())
MOD = 998244353

def cumprod_mod(arr):
    L = len(arr)
    Lsq = int(L**.5+1)
    arr = np.resize(arr,Lsq**2).reshape(Lsq,Lsq)
    for n in range(1,Lsq):
        arr[:,n] *= arr[:,n-1]; arr[:,n] %= MOD
    for n in range(1,Lsq):
        arr[n] *= arr[n-1,-1]; arr[n] %= MOD
    arr = arr.ravel()[:L]
    return arr

U = 2*10**6 + 100
x = np.arange(U,dtype=np.int64)
x[0] = 1
fact = cumprod_mod(x)

x = np.arange(U,0,-1,dtype=np.int64)
x[0] = pow(int(fact[-1]),MOD-2,MOD)
fact_inv = cumprod_mod(x)[::-1]

def comb(n,k):
    return fact[n] * fact_inv[k] % MOD * fact_inv[n-k] % MOD

x = comb(3*M+N-1,N-1) # 何でもいいから作る
y = N * comb(M+N-2,N-1) # どれかが 2M+1 以上

# 奇数がM+1個以上あるものを除く
odd_cnt = np.arange(M+2,3*M+1,2)
S = (3*M - odd_cnt)//2
bl = (odd_cnt <= N)&(S>=0)
odd_cnt = odd_cnt[bl]
S = S[bl]
arr = fact[N] * fact_inv[odd_cnt] % MOD * fact_inv[N-odd_cnt] % MOD
# SをN個で作る -> S+N-1 C N-1
arr *= fact[S+N-1] * fact_inv[N-1] % MOD * fact_inv[S] % MOD
arr %= MOD
z = arr.sum() % MOD

answer = x-y-z
answer %= MOD
print(answer)

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

N,M = map(int,input().split())
MOD = 998244353

U = 2 * 10 ** 6
fact = [1] * (U+1)
for n in range(1,U+1):
    fact[n] = fact[n-1] * n % MOD
fact_inv = [1] * (U+1)
fact_inv[U] = pow(fact[U],MOD-2,MOD)
for n in range(U,0,-1):
    fact_inv[n-1] = fact_inv[n] * n % MOD

def comb(n,k):
    if n == -1 and k == -1:
      return 1
    if n < 0:
      return 0
    if k < 0 or k > n:
      return 0
    return fact[n] * fact_inv[k] * fact_inv[n-k] % MOD

answer = comb(3*M+N-1,N-1) # 何でもいいから作る
answer -= N * comb(M+N-2,N-1) # どれかが 2M+1 以上

# 「ちょうど1」がM個より多くある場合を引かなければいけない
# （既に引いた場合と背反）
for k in range(M+1,3*M+1):
    sgn = -1 if (k-M)%2 else 1
    # k個決める、残りは和が3M-k, 個数N-k
    x = comb(N,k)
    x *= comb(3*M-k+N-k-1,N-k-1)
    answer += sgn * x * comb(k-1,k-1-M)
answer %= MOD

print(answer)

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

N,M = map(int,input().split())
MOD = 998244353

U = 2 * 10 ** 6
fact = [1] * (U+1)
for n in range(1,U+1):
    fact[n] = fact[n-1] * n % MOD
fact_inv = [1] * (U+1)
fact_inv[U] = pow(fact[U],MOD-2,MOD)
for n in range(U,0,-1):
    fact_inv[n-1] = fact_inv[n] * n % MOD

def comb(n,k):
    if n == -1 and k == -1:
      return 1
    if k < 0 or k > n:
      return 0
    return fact[n] * fact_inv[k] * fact_inv[n-k] % MOD

answer = comb(3*M+N-1,N-1) # 何でもいいから作る
answer -= N * comb(M+N-2,N-1) # どれかが 2M+1 以上

# 「ちょうど1」がM個より多くある場合を引かなければいけない
# （既に引いた場合と背反）
for k in range(M+1,3*M+1):
    sgn = -1 if (k-M)%2 else 1
    # k個決める、残りは和が3M-k, 個数N-k
    x = comb(N,k)
    x *= comb(3*M-k+N-k-1,N-k-1)
    answer += sgn * x * comb(k-1,k-1-M)
answer %= MOD

print(answer)

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

N,M = map(int,input().split())
MOD = 998244353

U = 2 * 10 ** 6
fact = [1] * (U+1)
for n in range(1,U+1):
    fact[n] = fact[n-1] * n % MOD
fact_inv = [1] * (U+1)
fact_inv[U] = pow(fact[U],MOD-2,MOD)
for n in range(U,0,-1):
    fact_inv[n-1] = fact_inv[n] * n % MOD

def comb(n,k):
    return fact[n] * fact_inv[k] * fact_inv[n-k] % MOD

answer = comb(3*M+N-1,N-1) # 何でもいいから作る
answer -= N * comb(M+N-2,N-1) # どれかが 2M+1 以上
answer %= MOD

print(answer)