Problem p02649 (python)

Fixed submission:
import sys
import numpy as np
from numba import njit

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

@njit('(i4[::1],i4,i8[::1])', cache=True)
def main(A, B, C):
    counts = np.zeros(1 << 18, np.int32)
    popcount = np.zeros(1 << B, np.int32)
    for i in range(B):
        popcount[1 << i:1 << i + 1] = popcount[:1 << i] + 1

    answer = 0
    for i in range(1 << B):
        k = popcount[i]
        t = 0
        for x in A & i:
            n = counts[x]
            counts[x] += 1
            t -= C[n]
            t += C[n + 1]
        for x in A & i:
            counts[x] = 0
        if k & 1:
            t = -t
        answer += t
    return answer

N, K, S, T = map(int, readline().split())
A = np.array(readline().split(), np.int32)

def convert_problem(S, T, A):
    ng = np.zeros(len(A), np.bool)
    B = np.zeros_like(A)
    n = 0
    for i in range(18):
        s, t = (S >> i) & 1, (T >> i) & 1
        if (s, t) == (0, 0):
            ng |= ((A >> i) & 1) == 1
        elif (s, t) == (1, 1):
            ng |= ((A >> i) & 1) == 0
        elif (s, t) == (1, 0):
            print(0)
            exit()
        else:
            B += ((A >> i) & 1) << n
            n += 1
    return B[~ng], n

A, B = convert_problem(S, T, A)

C = np.zeros((100, 100), np.int64)
C[0, 0] = 1
for n in range(1, 100):
    C[n, :-1] += C[n - 1, :-1]
    C[n, 1:] += C[n - 1, :-1]
C = C[:, 1:K + 1].sum(axis=1)

print(main(A, B, C))

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

@njit('(i4[::1],i4,i8[::1])', cache=True)
def main(A, B, C):
    counts = np.zeros(1 << 18, np.int32)
    popcount = np.zeros(1 << B, np.int32)
    for i in range(B):
        popcount[1 << i:1 << i + 1] = popcount[:1 << i] + 1

    answer = 0
    for i in range(1 << B):
        k = popcount[i]
        t = 0
        for x in A & i:
            n = counts[x]
            counts[x] += 1
            t -= C[n]
            t += C[n + 1]
        for x in A & i:
            counts[x] = 0
        if k & 1:
            t = -t
        answer += t
    return answer
  
if sys.argv[-1] == 'ONLINE_JUDGE':
    from numba.pycc import CC
    cc = CC('my_module')
    cc.export('main', '(i4[::1],i4,i8[::1])')(main)
    cc.compile()
    exit()
from my_module import main

N, K, S, T = map(int, readline().split())
A = np.array(readline().split(), np.int32)

def convert_problem(S, T, A):
    ng = np.zeros(len(A), np.bool)
    B = np.zeros_like(A)
    n = 0
    for i in range(18):
        s, t = (S >> i) & 1, (T >> i) & 1
        if (s, t) == (0, 0):
            ng |= ((A >> i) & 1) == 1
        elif (s, t) == (1, 1):
            ng |= ((A >> i) & 1) == 0
        elif (s, t) == (1, 0):
            print(0)
            exit()
        else:
            B += ((A >> i) & 1) << n
            n += 1
    return B[~ng], n

A, B = convert_problem(S, T, A)

C = np.zeros((100, 100), np.int64)
C[0, 0] = 1
for n in range(1, 100):
    C[n, :-1] += C[n - 1, :-1]
    C[n, 1:] += C[n - 1, :-1]
C = C[:, 1:K + 1].sum(axis=1)

print(main(A, B, C))

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

A = 1024

def main(N, V, W, v, L):
    def precompute(N, V, W):
        U = 10**5 + 10
        dp = np.zeros((A, U), np.int32)
        for v in range(1, min(A, N + 1)):
            p = v // 2
            dp[v] = dp[p]
            dp[v, W[v]:] = np.maximum(dp[v, W[v]:], dp[p, :-W[v]] + V[v])
            for i in range(1, U):
                dp[v][i] = max(dp[v][i], dp[v][i - 1])
        return dp

    dp = precompute(N, V, W)
    values = np.empty(A, np.int32)
    weights = np.empty(A, np.int32)

    def solve(i, lim, values, weights, dp):
        if i < A:
            return dp[i][lim]
        values[0] = 0
        weights[0] = 0
        p = 0
        while i >= A:
            for j in range(1 << p):
                values[j + (1 << p)] = values[j] + V[i]
                weights[j + (1 << p)] = weights[j] + W[i]
            p += 1
            i >>= 1
        best = 0
        for n in range(1 << p):
            if weights[n] > lim:
                continue
            x = dp[i][lim - weights[n]] + values[n]
            best = max(best, x)
        return best

    Q = len(v)
    for i in range(Q):
        x = solve(v[i], L[i], values, weights, dp)
        print(x)

        
if sys.argv[-1] == 'ONLINE_JUDGE':
    from numba.pycc import CC
    cc = CC('my_module')
    cc.export('main', '(i4,i4[:],i4[:],i4[:],i4[:])')(main)
    cc.compile()
    exit()

from my_module import main

N = int(readline())
stdin = np.array(read().split(), np.int32)
V = np.zeros(N + 1, np.int32)
W = np.zeros(N + 1, np.int32)
V[1:] = stdin[:N + N:2]
W[1:] = stdin[1:N + N:2]
query = stdin[N + N + 1:]
v, L = query[::2].copy(), query[1::2].copy()

main(N, V, W, v, L)

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

A = 512

def main(N, V, W, v, L):
    def precompute(N, V, W):
        U = 10**5 + 10
        dp = np.zeros((A, U), np.int32)
        for v in range(1, min(A, N + 1)):
            p = v // 2
            dp[v] = dp[p]
            dp[v, W[v]:] = np.maximum(dp[v, W[v]:], dp[p, :-W[v]] + V[v])
            for i in range(1, U):
                dp[v][i] = max(dp[v][i], dp[v][i - 1])
        return dp

    dp = precompute(N, V, W)
    values = np.empty(A, np.int32)
    weights = np.empty(A, np.int32)

    def solve(i, lim, values, weights, dp):
        if i < A:
            return dp[i][lim]
        values[0] = 0
        weights[0] = 0
        p = 0
        while i >= A:
            for j in range(1 << p):
                values[j + (1 << p)] = values[j] + V[i]
                weights[j + (1 << p)] = weights[j] + W[i]
            p += 1
            i >>= 1
        best = 0
        for n in range(1 << p):
            if weights[n] > lim:
                continue
            x = dp[i][lim - weights[n]] + values[n]
            best = max(best, x)
        return best

    Q = len(v)
    for i in range(Q):
        x = solve(v[i], L[i], values, weights, dp)
        print(x)

        
if sys.argv[-1] == 'ONLINE_JUDGE':
    from numba.pycc import CC
    cc = CC('my_module')
    cc.export('main', '(i4,i4[:],i4[:],i4[:],i4[:])')(main)
    cc.compile()
    exit()

from my_module import main

N = int(readline())
stdin = np.array(read().split(), np.int32)
V = np.zeros(N + 1, np.int32)
W = np.zeros(N + 1, np.int32)
V[1:] = stdin[:N + N:2]
W[1:] = stdin[1:N + N:2]
query = stdin[N + N + 1:]
v, L = query[::2], query[1::2]

main(N, V, W, v, L)