Problem p02647 (python)

Fixed submission:
import sys
import numpy as np
from numba import njit
 
readline = sys.stdin.readline
 
@njit('i8,i8,i8[::1]',cache=True)
def loop(n,k,a):
    for j in range(k):
        b = np.zeros(n+1, dtype=np.int64)
        for i in range(n):
            b[max(i-a[i], 0)] += 1
            b[min(i+a[i]+1, n)] -= 1
        a = np.cumsum(b)[:-1]
        if np.all(a==n):
            return(a)
            break
    else:
        return(a)
 
 
n, k = map(int, readline().split())
a = np.array(readline().split(), dtype=np.int64)
 
ans = loop(n,k,a)
print(*ans)

Buggy submission (verdict: Runtime Error):
import sys
import copy
from numba import njit

def I(): return int(sys.stdin.readline().rstrip())
def LI():return list(map(int,sys.stdin.readline().rstrip().split()))
def LS():return list(sys.stdin.readline().rstrip().split())

@njit('i8,i8,i8[::1]',cache=True)
def main(N,K,A):
    #print(N,K)
    #print(A)

    for k in range(K):
        B = [0]*N
        for i, a in enumerate(A):
            #print(i, a)
            if a == 0:
                B[i] += 1
            else:
                # i-(a+0.5) ~ i+(a+0.5)
                # i: 整数、a: ＋整数　なので +0.5 は切り捨てられる。
                #　影響範囲は、 i-a ~ i+a
                for p in range(i-a,i+a+1):
                    #print("p:",p)
                    if p >=0 and p < N:
                        B[p] += 1
                #print(B)
        if k < K-1:
            A = B
        else:
          break
    return B




if __name__ == '__main__':
    N,K = LI()[:]
    A = LI()
    B = main(N,K,A)
    print(B)

Buggy submission (verdict: Runtime Error):
import sys
import copy
from numba import njit

def I(): return int(sys.stdin.readline().rstrip())
def LI():return list(map(int,sys.stdin.readline().rstrip().split()))
def LS():return list(sys.stdin.readline().rstrip().split())

@njit('i8,i8,i8[::1]',cache=True)
def main(N,K,A):
    #print(N,K)
    #print(A)

    for k in range(K):
        B = [0]*N
        for i, a in enumerate(A):
            #print(i, a)
            if a == 0:
                B[i] += 1
            else:
                # i-(a+0.5) ~ i+(a+0.5)
                # i: 整数、a: ＋整数　なので +0.5 は切り捨てられる。
                #　影響範囲は、 i-a ~ i+a
                for p in range(i-a,i+a+1):
                    #print("p:",p)
                    if p >=0 and p < N:
                        B[p] += 1
                #print(B)
        if k < K-1:
            A = B
        else:
            return B




if __name__ == '__main__':
    N,K = LI()[:]
    A = LI()
    B = main(N,K,A)
    print(B)

Buggy submission (verdict: Runtime Error):
import sys
import copy
from numba import njit

def I(): return int(sys.stdin.readline().rstrip())
def LI():return list(map(int,sys.stdin.readline().rstrip().split()))
def LS():return list(sys.stdin.readline().rstrip().split())

@njit(parallel=True)
def main(N,K,A):
    #print(N,K)
    #print(A)

    for k in range(K):
        B = [0]*N
        for i, a in enumerate(A):
            #print(i, a)
            if a == 0:
                B[i] += 1
            else:
                # i-(a+0.5) ~ i+(a+0.5)
                # i: 整数、a: ＋整数　なので +0.5 は切り捨てられる。
                #　影響範囲は、 i-a ~ i+a
                for p in range(i-a,i+a+1):
                    #print("p:",p)
                    if p >=0 and p < N:
                        B[p] += 1
                #print(B)
        if k < K-1:
            A = B
        else:
            return B




if __name__ == '__main__':
    N,K = LI()[:]
    A = LI()
    B = main(N,K,A)
    print(B)