Problem p02585 (python)

Fixed submission:
import math
import sys
import os

sys.setrecursionlimit(10**7)

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

if os.getenv("LOCAL"):
    inputFile = basename_without_ext = os.path.splitext(os.path.basename(__file__))[0]+'.txt'
    sys.stdin = open(inputFile, "r")
INF = float("inf")

N,K = LI()
P = LI()
C = LI()

P = list(map(lambda x:x-1,P))

res = -INF
used = [False] * N
ss = []

for i in range(N):
    # print('i',i)
    if used[i]:
        continue
    cur = i
    s = []
    while not used[cur]:
        used[cur] = True
        s.append(C[cur])
        cur = P[cur]
        # print(cur)
    ss.append(s)

# print(ss)

for vec in ss:
    M = len(vec)
    # 2周の累積和
    sum = [0]*(2*M + 1)
    for i in range(2*M):
        sum[i+1] = sum[i]+vec[i%M]
    
    # amari[r] := 連続するr個の最大値
    amari = [-INF]*M
    for i in range(M):
        for j in range(M):
            amari[j] = max(amari[j],sum[i+j] - sum[i])
    
    for r in range(M):
        if r>K:
            continue
        q = (K-r)//M
        if r==0 and q==0:
            continue
        if sum[M]>0:
            res = max(res, amari[r] + sum[M] * q)
        elif r > 0:
            res = max(res, amari[r])
print(res)

Buggy submission (verdict: Runtime Error):
import bisect,collections,copy,heapq,itertools,math,string
from collections import defaultdict as D
from functools import reduce
import numpy as np
import sys
import os
from operator import mul
from numba import njit, b1, i4, i8, f8

sys.setrecursionlimit(10**7)

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

if os.getenv("LOCAL"):
    inputFile = basename_without_ext = os.path.splitext(os.path.basename(__file__))[0]+'.txt'
    sys.stdin = open(inputFile, "r")
INF = float("inf")
IINF = 10 ** 18
MOD = 10 ** 9 + 7
# MOD = 998244353    

N,K = LI()
P = LI()
C = LI()
C = np.array(C,dtype='int')
P = list(map(lambda x:x-1,P))
P = np.array(P,dtype='int')

# def rec(visited,next,s):
#     if next in s:
#         return visited,visited.index(next)
#     else:
#         s.add(next)
#         visited.append(next)
#         return rec(visited,P[next],s)

@njit((i8, i8[:], i8[:]), cache=True)
def solve(K,P,C):
    N = len(P)
    # INF = 1 << 60
    best_score = -INF
    for i in range(N):
        s = set([i])
        visited = [i]
        a = 0
        # visited,loopStart = rec(visited,P[i],s)
        next = P[i]
        while True:
            if next in s:
                loopStart = visited.index(next)
                break
            else:
                s.add(next)
                visited.append(next)
                next = P[next]

        # print(visited,loopStart)
        al = np.zeros(len(visited))
        # al = np.array([],dtype='int')
        if K < len(visited):
            for j in range(len(visited)-1):
                # al = np.append(al,C[visited[j+1]])
                al[j+1]=C[visited[j+1]]
            # print(al)
            # a_cum = np.array([],dtype='int')
            a_cum = np.zeros(len(al))
            a_cum = al.cumsum()
            a = a_cum.max()
            # print(a)
            # a = acum.max()
        else:        
            for j in range(len(visited)-1):
                al = np.append(al,C[visited[j+1]])
            visitedLoop = visited[loopStart:]
            # print("visitedLoop: ",visitedLoop)
            sum_visitedLoop = 0
            for j in range(len(visitedLoop)):
                sum_visitedLoop += C[visitedLoop[j]]
            if sum_visitedLoop<=0:
                a = al.cumsum().max()
            else:
                # 残り回数を循環配列の長さで割った余り
                q,mod = divmod((K-len(visited)),len(visitedLoop))
                # print('q,mod:',q,mod)
                # am = np.array([0],dtype='int')
                am = np.zeros(mod)

                for j in range(mod+1):
                    # am = np.append(am,C[visitedLoop[j]])
                    am[j] = C[visitedLoop[j]]
                    am_sum = am.cumsum().max()
                # print(am)
                # print('al:',al)
                # print(al.cumsum()[-1],q,sum_visitedLoop, am_sum)
                a = al.cumsum()[-1] + q*sum_visitedLoop + am_sum
        best_score = max(best_score, a)
    return best_score

res = solve(K,P,C)
print(res)

Buggy submission (verdict: Runtime Error):
import bisect,collections,copy,heapq,itertools,math,string
from collections import defaultdict as D
from functools import reduce
import numpy as np
import sys
import os
from operator import mul
from numba import njit

sys.setrecursionlimit(10**7)

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

if os.getenv("LOCAL"):
    inputFile = basename_without_ext = os.path.splitext(os.path.basename(__file__))[0]+'.txt'
    sys.stdin = open(inputFile, "r")
INF = float("inf")
IINF = 10 ** 18
MOD = 10 ** 9 + 7
# MOD = 998244353    

N,K = LI()
P = LI()
C = LI()
C = np.array(C,dtype='int')
P = list(map(lambda x:x-1,P))
P = np.array(P,dtype='int')

# def rec(visited,next,s):
#     if next in s:
#         return visited,visited.index(next)
#     else:
#         s.add(next)
#         visited.append(next)
#         return rec(visited,P[next],s)

@njit
def solve(N,K,P,C):
    ans = [0] * N
    for i in range(N):
        s = set([i])
        visited = [i]
        a = 0
        # visited,loopStart = rec(visited,P[i],s)
        next = P[i]
        while True:
            if next in s:
                loopStart = visited.index(next)
                break
            else:
                s.add(next)
                visited.append(next)
                next = P[next]

        # print(visited,loopStart)
        al = np.zeros(len(visited))
        # al = np.array([],dtype='int')
        if K < len(visited):
            for j in range(len(visited)-1):
                # al = np.append(al,C[visited[j+1]])
                al[j+1]=C[visited[j+1]]
            # print(al)
            # a_cum = np.array([],dtype='int')
            a_cum = np.zeros(len(al))
            a_cum = al.cumsum()
            a = a_cum.max()
            # print(a)
            # a = acum.max()
        else:        
            for j in range(len(visited)-1):
                al = np.append(al,C[visited[j+1]])
            visitedLoop = visited[loopStart:]
            # print("visitedLoop: ",visitedLoop)
            sum_visitedLoop = 0
            for j in range(len(visitedLoop)):
                sum_visitedLoop += C[visitedLoop[j]]
            if sum_visitedLoop<=0:
                a = al.cumsum().max()
            else:
                # 残り回数を循環配列の長さで割った余り
                q,mod = divmod((K-len(visited)),len(visitedLoop))
                # print('q,mod:',q,mod)
                # am = np.array([0],dtype='int')
                am = np.zeros(mod)

                for j in range(mod+1):
                    # am = np.append(am,C[visitedLoop[j]])
                    am[j] = C[visitedLoop[j]]
                    am_sum = am.cumsum().max()
                # print(am)
                # print('al:',al)
                # print(al.cumsum()[-1],q,sum_visitedLoop, am_sum)
                a = al.cumsum()[-1] + q*sum_visitedLoop + am_sum
        ans[i]=a
    return ans
res = solve(N,K,P,C)
print(max(res))

Buggy submission (verdict: Runtime Error):
import bisect,collections,copy,heapq,itertools,math,string
from collections import defaultdict as D
from functools import reduce
import numpy as np
import sys
import os
from operator import mul
from numba import njit

sys.setrecursionlimit(10**7)

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

if os.getenv("LOCAL"):
    inputFile = basename_without_ext = os.path.splitext(os.path.basename(__file__))[0]+'.txt'
    sys.stdin = open(inputFile, "r")
INF = float("inf")
IINF = 10 ** 18
MOD = 10 ** 9 + 7
# MOD = 998244353    

N,K = LI()
P = LI()
C = LI()

ans = [0] * N
P = list(map(lambda x:x-1,P))

def rec(visited,next,s):
    if next in s:
        return visited,visited.index(next)
    else:
        s.add(next)
        visited.append(next)
        return rec(visited,P[next],s)

# @njit(nopython=True)
def solve():
    for i in range(N):
        s = set([i])
        visited = [i]
        a = 0
        visited,loopStart = rec(visited,P[i],s)
        # print(visited,loopStart)
        al = np.array([],dtype='int')
        if K < len(visited):
            for j in range(len(visited)-1):
                al = np.append(al,C[visited[j+1]])
            # print(al)
            a = al[1:].cumsum().max()
            # print(a)
            # a = acum.max()
        else:        
            for j in range(len(visited)-1):
                al = np.append(al,C[visited[j+1]])
            visitedLoop = visited[loopStart:]
            # print("visitedLoop: ",visitedLoop)
            sum_visitedLoop = 0
            for j in range(len(visitedLoop)):
                sum_visitedLoop += C[visitedLoop[j]]
            if sum_visitedLoop<=0:
                a = al.cumsum().max()
            else:
                # 残り回数を循環配列の長さで割った余り
                q,mod = divmod((K-len(visited)),len(visitedLoop))
                # print('q,mod:',q,mod)
                am = np.array([0],dtype='int')
                for j in range(mod+1):
                    am = np.append(am,C[visitedLoop[j]])
                    am_sum = am.cumsum().max()
                # print(am)
                # print('al:',al)
                # print(al.cumsum()[-1],q,sum_visitedLoop, am_sum)
                a = al.cumsum()[-1] + q*sum_visitedLoop + am_sum
        ans[i]=a
solve()
print(max(ans))

Buggy submission (verdict: Runtime Error):
import bisect,collections,copy,heapq,itertools,math,string
from collections import defaultdict as D
from functools import reduce
import numpy as np
import sys
import os
from operator import mul

sys.setrecursionlimit(10**7)

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

if os.getenv("LOCAL"):
    inputFile = basename_without_ext = os.path.splitext(os.path.basename(__file__))[0]+'.txt'
    sys.stdin = open(inputFile, "r")
INF = float("inf")
IINF = 10 ** 18
MOD = 10 ** 9 + 7
# MOD = 998244353    

N,K = LI()
P = LI()
C = LI()

ans = [0] * N
P = list(map(lambda x:x-1,P))

def rec(visited,next):
    if next in s:
        return visited,visited.index(next)
    else:
        s.add(next)
        visited.append(next)
        return rec(visited,P[next])

for i in range(N):
    s = set([i])
    visited = [i]
    a = 0
    visited,loopStart = rec(visited,P[i])
    # print(visited,loopStart)
    al = np.array([],dtype='int')
    if K < len(visited):
        for j in range(len(visited)-1):
            al = np.append(al,C[visited[j+1]])
        # print(al)
        a = al[1:].cumsum().max()
        # print(a)
        # a = acum.max()
    else:        
        for j in range(len(visited)-1):
            al = np.append(al,C[visited[j+1]])
        visitedLoop = visited[loopStart:]
        # print("visitedLoop: ",visitedLoop)
        sum_visitedLoop = 0
        for j in range(len(visitedLoop)):
            sum_visitedLoop += C[visitedLoop[j]]
        if sum_visitedLoop<=0:
            a = al[1:].cumsum().max()
        else:
            # 残り回数を循環配列の長さで割った余り
            q,mod = divmod((K-len(visited)),len(visitedLoop))
            # print('q,mod:',q,mod)
            am = np.array([0],dtype='int')
            for j in range(mod+1):
                am = np.append(am,C[visitedLoop[j]])
                am_sum = am.cumsum().max()
            # print(am)
            # print('al:',al)
            # print(al.cumsum()[-1],q,sum_visitedLoop, am_sum)
            a = al.cumsum()[-1] + q*sum_visitedLoop + am_sum
    ans[i]=a

print(max(ans))