Problem p02539 (python)

Fixed submission:
import os
import sys
from heapq import heappop, heappush

import numpy as np


# https://github.com/atcoder/ac-library/blob/master/atcoder/convolution.hpp

def solve(inp):
    def ceil_pow2(n):
        x = 0
        while 1 << x < n:
            x += 1
        return x

    def bit_scan_forward(n):
        x = 0
        while n & 1 == 0:
            n >>= 1
            x += 1
        return x

    def pow_mod(x, n, m):
        r = 1
        y = x % m
        while n:
            if n & 1:
                r = (r * y) % m
            y = y * y % m
            n >>= 1
        return r

    def get_primitive_root(m):
        if m == 2: return 1
        if m == 167772161: return 3
        if m == 469762049: return 3
        if m == 754974721: return 11
        if m == 998244353: return 3
        divs = [2]
        x = (m - 1) // 2
        while x & 1 == 0:
            x >>= 1
        i = 3
        while i * i <= x:
            if x % i == 0:
                divs.append(i)
                x //= i
                while x % i == 0:
                    x //= i
            i += 2
        if x > 1:
            divs.append(x)
        g = 2
        while True:
            ok = True
            for d in divs:
                if pow_mod(g, (m - 1) // d, m) == 1:
                    ok = False
                    break
            if ok:
                return g

    def butterfly_prepare(mod, primitive_root):
        sum_e = np.zeros(30, np.int64)
        sum_ie = np.zeros(30, np.int64)
        es = np.zeros(30, np.int64)
        ies = np.zeros(30, np.int64)
        cnt2 = bit_scan_forward(mod - 1)
        e = pow_mod(primitive_root, (mod - 1) >> cnt2, mod)
        ie = pow_mod(e, mod - 2, mod)
        for i in range(cnt2, 1, -1):
            es[i - 2] = e
            ies[i - 2] = ie
            e = e * e % mod
            ie = ie * ie % mod
        now_e = 1
        now_ie = 1
        for i in range(cnt2 - 1):
            sum_e[i] = es[i] * now_e % mod
            sum_ie[i] = ies[i] * now_ie % mod
            now_e = now_e * ies[i] % mod
            now_ie = now_ie * es[i] % mod
        return sum_e, sum_ie

    def butterfly(aaa, mod, sum_e):
        n = aaa.size
        h = ceil_pow2(n)

        for ph in range(1, h + 1):
            w = 1 << (ph - 1)
            p = 1 << (h - ph)
            now = 1
            for s in range(w):
                offset = s << (h - ph + 1)
                for i in range(p):
                    l = aaa[i + offset]
                    r = aaa[i + offset + p] * now % mod
                    aaa[i + offset] = (l + r) % mod
                    aaa[i + offset + p] = (l - r) % mod
                now = now * sum_e[bit_scan_forward(~s)] % mod

    def butterfly_inv(aaa, mod, sum_ie):
        n = aaa.size
        h = ceil_pow2(n)

        for ph in range(h, 0, -1):
            w = 1 << (ph - 1)
            p = 1 << (h - ph)
            inow = 1
            for s in range(w):
                offset = s << (h - ph + 1)
                for i in range(p):
                    l = aaa[i + offset]
                    r = aaa[i + offset + p]
                    aaa[i + offset] = (l + r) % mod
                    aaa[i + offset + p] = ((l - r) * inow) % mod
                inow = inow * sum_ie[bit_scan_forward(~s)] % mod

    MOD = 998244353
    primitive_root = get_primitive_root(MOD)
    sum_e, sum_ie = butterfly_prepare(MOD, primitive_root)

    def convolution(aaa, bbb, MOD):
        n = aaa.size
        m = bbb.size
        z = 1 << ceil_pow2(n + m - 1)
        raaa = np.zeros(z, np.int64)
        rbbb = np.zeros(z, np.int64)
        raaa[:n] = aaa
        rbbb[:m] = bbb
        butterfly(raaa, MOD, sum_e)
        butterfly(rbbb, MOD, sum_e)
        ccc = raaa * rbbb % MOD
        butterfly_inv(ccc, MOD, sum_ie)
        iz = pow_mod(z, MOD - 2, MOD)
        result = ccc[:n + m - 1] * iz % MOD
        return result

    def mod_pow(x, a, MOD):
        ret = 1
        cur = x
        while a:
            if a & 1:
                ret = ret * cur % MOD
            cur = cur * cur % MOD
            a >>= 1
        return ret

    def precompute_factorials(n, MOD):
        factorials = np.ones(n + 1, dtype=np.int64)
        for m in range(2, n + 1):
            factorials[m] = factorials[m - 1] * m % MOD
        inversions = np.ones(n + 1, dtype=np.int64)
        inversions[n] = mod_pow(factorials[n], MOD - 2, MOD)
        for m in range(n, 2, -1):
            inversions[m - 1] = inversions[m] * m % MOD
        return factorials, inversions

    def counter(arr):
        cnt = {0: 0}
        for a in arr:
            if a in cnt:
                cnt[a] += 1
            else:
                cnt[a] = 1
        if cnt[0] == 0:
            del cnt[0]
        return cnt

    n = inp[0]
    hhh = inp[1:]
    hhh_cnt = counter(hhh)

    same_heights = []
    max_pair_cnt = 0
    max_same_height = 0
    for c in hhh_cnt.values():
        if c <= 1:
            continue
        max_pair_cnt += c // 2
        max_same_height = max(max_same_height, c)
        same_heights.append(c)

    pairs = np.zeros(n + 1, np.int64)
    pairs[0] = 1
    for i in range(1, n + 1):
        pairs[i] = pairs[i - 1] * (i * 2 - 1) % MOD

    facts, finvs = precompute_factorials(max_same_height, MOD)

    same_heights.sort()
    pc = -1
    cmbs = [np.ones(1, np.int64)]
    q = [(1, 0)]
    for i, c in enumerate(same_heights):
        if pc == c:
            cmb = cmbs[-1].copy()
        else:
            cmb = np.zeros(c // 2 + 1, np.int64)
            for j in range(c // 2 + 1):
                cmb[j] = facts[c] * finvs[j * 2] % MOD * finvs[c - j * 2] % MOD * pairs[j] % MOD
            pc = c
        cmbs.append(cmb)
        q.append((cmb.size, i + 1))

    while len(q) > 1:
        _, i = heappop(q)
        _, j = heappop(q)
        cmbs[i] = convolution(cmbs[i], cmbs[j], MOD)
        heappush(q, (cmbs[i].size, i))

    result = cmbs[q[0][1]]

    ans = 0
    for i, v in enumerate(result):
        if i & 1:
            ans = (ans - v * pairs[n - i]) % MOD
        else:
            ans = (ans + v * pairs[n - i]) % MOD

    return ans


if sys.argv[-1] == 'ONLINE_JUDGE':
    from numba.pycc import CC

    cc = CC('my_module')
    cc.export('solve', '(i8[:],)')(solve)
    cc.compile()
    exit()

if os.name == 'posix':
    # noinspection PyUnresolvedReferences
    from my_module import solve
else:
    from numba import njit

    solve = njit('(i8[:],)', cache=True)(solve)
    print('compiled', file=sys.stderr)

inp = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=' ')
ans = solve(inp)
print(ans)

Buggy submission (verdict: Runtime Error):
import os
import sys
from heapq import heappop, heappush

import numpy as np


# https://github.com/atcoder/ac-library/blob/master/atcoder/convolution.hpp

def solve(inp):
    def ceil_pow2(n):
        x = 0
        while 1 << x < n:
            x += 1
        return x

    def bit_scan_forward(n):
        x = 0
        while n & 1 == 0:
            n >>= 1
            x += 1
        return x

    def pow_mod(x, n, m):
        r = 1
        y = x % m
        while n:
            if n & 1:
                r = (r * y) % m
            y = y * y % m
            n >>= 1
        return r

    def get_primitive_root(m):
        if m == 2: return 1
        if m == 167772161: return 3
        if m == 469762049: return 3
        if m == 754974721: return 11
        if m == 998244353: return 3
        divs = [2]
        x = (m - 1) // 2
        while x & 1 == 0:
            x >>= 1
        i = 3
        while i * i <= x:
            if x % i == 0:
                divs.append(i)
                x //= i
                while x % i == 0:
                    x //= i
            i += 2
        if x > 1:
            divs.append(x)
        g = 2
        while True:
            ok = True
            for d in divs:
                if pow_mod(g, (m - 1) // d, m) == 1:
                    ok = False
                    break
            if ok:
                return g

    def butterfly_prepare(mod, primitive_root):
        sum_e = np.zeros(30, np.int64)
        sum_ie = np.zeros(30, np.int64)
        es = np.zeros(30, np.int64)
        ies = np.zeros(30, np.int64)
        cnt2 = bit_scan_forward(mod - 1)
        e = pow_mod(primitive_root, (mod - 1) >> cnt2, mod)
        ie = pow_mod(e, mod - 2, mod)
        for i in range(cnt2, 1, -1):
            es[i - 2] = e
            ies[i - 2] = ie
            e = e * e % mod
            ie = ie * ie % mod
        now_e = 1
        now_ie = 1
        for i in range(cnt2 - 1):
            sum_e[i] = es[i] * now_e % mod
            sum_ie[i] = ies[i] * now_ie % mod
            now_e = now_e * ies[i] % mod
            now_ie = now_ie * es[i] % mod
        return sum_e, sum_ie

    def butterfly(aaa, mod, sum_e):
        n = aaa.size
        h = ceil_pow2(n)

        for ph in range(1, h + 1):
            w = 1 << (ph - 1)
            p = 1 << (h - ph)
            now = 1
            for s in range(w):
                offset = s << (h - ph + 1)
                for i in range(p):
                    l = aaa[i + offset]
                    r = aaa[i + offset + p] * now % mod
                    aaa[i + offset] = (l + r) % mod
                    aaa[i + offset + p] = (l - r) % mod
                now = now * sum_e[bit_scan_forward(~s)] % mod

    def butterfly_inv(aaa, mod, sum_ie):
        n = aaa.size
        h = ceil_pow2(n)

        for ph in range(h, 0, -1):
            w = 1 << (ph - 1)
            p = 1 << (h - ph)
            inow = 1
            for s in range(w):
                offset = s << (h - ph + 1)
                for i in range(p):
                    l = aaa[i + offset]
                    r = aaa[i + offset + p]
                    aaa[i + offset] = (l + r) % mod
                    aaa[i + offset + p] = ((l - r) * inow) % mod
                inow = inow * sum_ie[bit_scan_forward(~s)] % mod

    MOD = 998244353
    primitive_root = get_primitive_root(MOD)
    sum_e, sum_ie = butterfly_prepare(MOD, primitive_root)

    def convolution(aaa, bbb, MOD):
        n = aaa.size
        m = bbb.size
        z = 1 << ceil_pow2(n + m - 1)
        raaa = np.zeros(z, np.int64)
        rbbb = np.zeros(z, np.int64)
        raaa[:n] = aaa
        rbbb[:m] = bbb
        butterfly(raaa, MOD, sum_e)
        butterfly(rbbb, MOD, sum_e)
        ccc = raaa * rbbb % MOD
        butterfly_inv(ccc, MOD, sum_ie)
        iz = pow_mod(z, MOD - 2, MOD)
        result = ccc[:n + m - 1] * iz % MOD
        return result

    def mod_pow(x, a, MOD):
        ret = 1
        cur = x
        while a:
            if a & 1:
                ret = ret * cur % MOD
            cur = cur * cur % MOD
            a >>= 1
        return ret

    def precompute_factorials(n, MOD):
        factorials = np.ones(n + 1, dtype=np.int64)
        for m in range(2, n + 1):
            factorials[m] = factorials[m - 1] * m % MOD
        inversions = np.ones(n + 1, dtype=np.int64)
        inversions[n] = mod_pow(factorials[n], MOD - 2, MOD)
        for m in range(n, 2, -1):
            inversions[m - 1] = inversions[m] * m % MOD
        return factorials, inversions

    def counter(arr):
        cnt = {0: 0}
        for a in arr:
            if a in cnt:
                cnt[a] += 1
            else:
                cnt[a] = 1
        if cnt[0] == 0:
            del cnt[0]
        return cnt

    n = inp[0]
    hhh = inp[1:]
    hhh_cnt = counter(hhh)
    same_heights = []
    max_pair_cnt = 0
    max_same_height = 0
    for c in hhh_cnt.values():
        if c <= 1:
            continue
        max_pair_cnt += c // 2
        max_same_height = max(max_same_height, c)
        same_heights.append(c)

    pairs = np.zeros(n + 1, np.int64)
    pairs[0] = 1
    for i in range(1, n + 1):
        pairs[i] = pairs[i - 1] * (i * 2 - 1) % MOD

    facts, finvs = precompute_factorials(max_same_height, MOD)

    same_heights.sort()
    pc = -1
    cmbs = []
    q = []
    for i, c in enumerate(same_heights):
        if pc == c:
            cmb = cmbs[-1].copy()
        else:
            cmb = np.zeros(c // 2 + 1, np.int64)
            for j in range(c // 2 + 1):
                cmb[j] = facts[c] * finvs[j * 2] % MOD * finvs[c - j * 2] % MOD * pairs[j] % MOD
            pc = c
        cmbs.append(cmb)
        q.append((cmb.size, i))

    while len(q) > 1:
        _, i = heappop(q)
        _, j = heappop(q)
        cmbs[i] = convolution(cmbs[i], cmbs[j], MOD)
        heappush(q, (cmbs[i].size, i))

    result = cmbs[q[0][1]]

    ans = 0
    for i, v in enumerate(result):
        if i & 1:
            ans = (ans - v * pairs[n - i]) % MOD
        else:
            ans = (ans + v * pairs[n - i]) % MOD

    return ans


if sys.argv[-1] == 'ONLINE_JUDGE':
    from numba.pycc import CC

    cc = CC('my_module')
    cc.export('solve', '(i8[:],)')(solve)
    cc.compile()
    exit()

if os.name == 'posix':
    # noinspection PyUnresolvedReferences
    from my_module import solve
else:
    from numba import njit

    solve = njit('(i8[:],)', cache=True)(solve)
    print('compiled', file=sys.stderr)

inp = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=' ')
ans = solve(inp)
print(ans)