Problem p02585 (python)

Fixed submission:
(n, k), p, c = [[*map(int, i.split())] for i in open(0)]

def solve(x, t):
    visit = [0] * n
    visit[x] = 1
    loop = [0] * n
    count = 0
    ans = c[p[x] - 1]
    l = True
    sub = 0
    while True:
        x = p[x] - 1
        if l:
            if visit[x]:
                if loop[x]:
                    ln = sum(loop)
                    if t > l:
                        sub += (t//ln -1)*count*(count>0)
                        t %= ln
                        t += ln
                    l = False
                count += c[x]
                loop[x] = 1
            visit[x] = 1
        sub += c[x]
        t -= 1
        ans = max(sub, ans)
        if t < 1:
            return ans


print(max(solve(i, k) for i in range(n)))

Buggy submission (verdict: Runtime Error):
#import sys

import numpy as np


def main(n,k,p,c):
    ans = -10**9
    for i in range(n):
        x = i
        t = k
        count = 0
        l = True
        sub = 0
        ln = 0
        while True:
            x = p[x] - 1
            if l:
                count += c[x]
                ln += 1
                if x == i:
                    if t > ln:
                        if count > 0:
                            sub += (t // ln - 1) * count
                            t %= ln
                            t += ln
                        else:
                            t = 0
                    l = False
            sub += c[x]
            t -= 1
            ans = max(sub, ans)
            if t < 1:
                break
    return ans


# >>> numba compile >>>
#if sys.argv[-1]=='ONLINE_JUDGE':
#  from numba.pycc import CC
#  cc=CC('my_module')
#  cc.export('main','i8(i8,i8,i8[:],i8[:])')(main)
#  cc.compile()
#  exit()
#from my_module import main
# <<< numba compile <<<
(n, k), p, c = [np.int_(i.split()) for i in open(0)]
print(main(n,k,p,c))

Buggy submission (verdict: Runtime Error):
import sys

import numpy as np


def main(n,k,p,c):
    ans = -10**9
    for i in range(n):
        x = i
        t = k
        count = 0
        l = True
        sub = 0
        ln = 0
        while True:
            x = p[x] - 1
            if l:
                count += c[x]
                ln += 1
                if x == i:
                    if t > ln:
                        if count > 0:
                            sub += (t // ln - 1) * count
                            t %= ln
                            t += ln
                        else:
                            t = 0
                    l = False
            sub += c[x]
            t -= 1
            ans = max(sub, ans)
            if t < 1:
                break
    return ans


# >>> numba compile >>>
#if sys.argv[-1]=='ONLINE_JUDGE':
#  from numba.pycc import CC
#  cc=CC('my_module')
#  cc.export('main','i8(i8,i8,i8[:],i8[:])')(main)
#  cc.compile()
#  exit()
#from my_module import main
# <<< numba compile <<<
(n, k), p, c = [np.int_(i.split()) for i in open(0)]
print(main(n,k,p,c))

Buggy submission (verdict: Runtime Error):
import os
import sys

import numpy as np


def solve(n,k,p,c):
    ans = -10**9
    for i in range(n):
        x = i
        t = k
        count = 0
        l = True
        sub = 0
        ln = 0
        while True:
            x = p[x] - 1
            if l:
                count += c[x]
                ln += 1
                if x == i:
                    if t > ln:
                        if count > 0:
                            sub += (t // ln - 1) * count
                            t %= ln
                            t += ln
                        else:
                            t = 0
                    l = False
            sub += c[x]
            t -= 1
            ans = max(sub, ans)
            if t < 1:
                break
    return ans


# >>> numba compile >>>
if sys.argv[-1]=='ONLINE_JUDGE':
  from numba.pycc import CC
  cc=CC('my_module')
  cc.export('solve','i8(i8,i8,i8[:],i8[:])')(main)
  cc.compile()
  exit()
from my_module import main
# <<< numba compile <<<
(n, k), p, c = [np.int_(i.split()) for i in open(0)]
print(solve(n,k,p,c))

Buggy submission (verdict: Runtime Error):
import os
import sys

import numpy as np


def solve(x, t):
    visit = [0] * n
    visit[x] = 1
    loop = [0] * n
    count = 0
    ans = c[p[x] - 1]
    l = True
    sub = 0
    while True:
        x = p[x] - 1
        if l:
            if visit[x]:
                if loop[x]:
                    ln = sum(loop)
                    if t > l:
                        sub += (t // ln - 1) * count * (count > 0)
                        t %= ln
                        t += ln
                    l = False
                count += c[x]
                loop[x] = 1
            visit[x] = 1
        sub += c[x]
        t -= 1
        ans = max(sub, ans)
        if t < 1:
            return ans


# >>> numba compile >>>
numba_config = [
    [solve, "i8(i8,i8)"],
]
if sys.argv[-1] == "ONLINE_JUDGE":
    from numba import njit
    from numba.pycc import CC

    cc = CC("my_module")
    for func, signature in numba_config:
        vars()[func.__name__] = njit(signature)(func)
        cc.export(func.__name__, signature)(func)
    cc.compile()
    exit()
elif os.name == "posix":
    exec(
        f"from my_module import {','.join(func.__name__ for func, _ in numba_config)}"
    )
else:
    from numba import njit

    for func, signature in numba_config:
        vars()[func.__name__] = njit(signature, cache=True)(func)
    print("compiled!", file=sys.stderr)
# <<< numba compile <<<
(n, k), p, c = [np.int_(i.split()) for i in open(0)]
print(max(solve(i, k) for i in range(n)))