Problem p02609 (python)

Fixed submission:
import sys
import numpy as np
import numba
from numba import njit
i8 = numba.int64

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

N = int(readline())
X = np.array(list(read().rstrip()), np.int64) - ord('0')

@njit((i8, ), cache=True)
def popcount(n):
    ret = 0
    while n:
        ret += n & 1
        n >>= 1
    return ret


@njit((i8, ), cache=True)
def f(n):
    t = 0
    while n:
        t += 1
        n %= popcount(n)
    return t

@njit(cache=True)
def main(X):

    popX = X.sum()
    MOD1 = popX + 1
    MOD2 = popX - 1
    if popX == 1:
        MOD2 = 1
    pow1 = np.ones_like(X)
    pow2 = np.ones_like(X)
    for n in range(1, len(X)):
        pow1[n] = 2 * pow1[n - 1] % MOD1
        pow2[n] = 2 * pow2[n - 1] % MOD2

    X1, X2 = 0, 0
    for x in X:
        X1 = (2 * X1 + x) % MOD1
        X2 = (2 * X2 + x) % MOD2

    ans = np.empty(N, np.int64)
    for i in range(N):
        x = X[N - 1 - i]
        if x == 0:
            nxt = (X1 + pow1[i]) % (popX + 1)
            ans[i] = 1 + f(nxt)
        elif x == 1:
            if popX == 1:
                ans[i] = 0
            else:
                nxt = (X2 - pow2[i]) % (popX - 1)
                ans[i] = 1 + f(nxt)
    return ans[::-1]

ans = main(X)
print('\n'.join(map(str, ans.tolist())))

Buggy submission (verdict: Runtime Error):
import sys
import numpy as np
import numba
from numba import njit
i8 = numba.int64

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

@njit((i8, ), cache=True)
def popcount(n):
    ret = 0
    while n:
        ret += n & 1
        n >>= 1
    return ret


@njit((i8, ), cache=True)
def f(n):
    t = 0
    while n:
        t += 1
        n %= popcount(n)
    return t

@njit((i8[:],), cache=True)
def main(X):

    popX = X.sum()
    MOD1 = popX + 1
    MOD2 = popX - 1
    if popX == 1:
        MOD2 = 1
    pow1 = np.ones_like(X)
    pow2 = np.ones_like(X)
    for n in range(1, len(X)):
        pow1[n] = 2 * pow1[n - 1] % MOD1
        pow2[n] = 2 * pow2[n - 1] % MOD2

    X1, X2 = 0, 0
    for x in X:
        X1 = (2 * X1 + x) % MOD1
        X2 = (2 * X2 + x) % MOD2

    ans = np.empty(N, np.int64)
    for i in range(N):
        x = X[N - 1 - i]
        if x == 0:
            nxt = (X1 + pow1[i]) % (popX + 1)
            ans[i] = 1 + f(nxt)
        elif x == 1:
            if popX == 1:
                ans[i] = 0
            else:
                nxt = (X2 - pow2[i]) % (popX - 1)
                ans[i] = 1 + f(nxt)
    return ans[::-1]

N = int(readline())
X = np.array(list(read().rstrip()), np.int64) - ord('0')
ans = main(X)
print('\n'.join(map(str, ans.tolist())))

Buggy submission (verdict: Runtime Error):
import sys
import numpy as np
import numba
from numba import njit
i8 = numba.int64

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

@njit((i8, ), cache=True)
def popcount(n):
    ret = 0
    while n:
        ret += n & 1
        n >>= 1
    return ret


@njit((i8, ), cache=True)
def f(n):
    t = 0
    while n:
        t += 1
        n %= popcount(n)
    return t

@njit((i8[:]), cache=True)
def main(X):

    popX = X.sum()
    MOD1 = popX + 1
    MOD2 = popX - 1
    if popX == 1:
        MOD2 = 1
    pow1 = np.ones_like(X)
    pow2 = np.ones_like(X)
    for n in range(1, len(X)):
        pow1[n] = 2 * pow1[n - 1] % MOD1
        pow2[n] = 2 * pow2[n - 1] % MOD2

    X1, X2 = 0, 0
    for x in X:
        X1 = (2 * X1 + x) % MOD1
        X2 = (2 * X2 + x) % MOD2

    ans = np.empty(N, np.int64)
    for i in range(N):
        x = X[N - 1 - i]
        if x == 0:
            nxt = (X1 + pow1[i]) % (popX + 1)
            ans[i] = 1 + f(nxt)
        elif x == 1:
            if popX == 1:
                ans[i] = 0
            else:
                nxt = (X2 - pow2[i]) % (popX - 1)
                ans[i] = 1 + f(nxt)
    return ans[::-1]

N = int(readline())
X = np.array(list(read().rstrip()), np.int64) - ord('0')
ans = main(X)
print('\n'.join(map(str, ans.tolist())))