Problem p02585 (python)

Fixed submission:
import bisect
import sys,io
import math
from collections import deque
from heapq import heappush, heappop

def sRaw():
    return input().rstrip("\r")


def iRaw():
    return int(input())


def ssRaw():
    return input().split()


def isRaw():
    return list(map(int, ssRaw()))


INF = 1 << 29

DIV = (10**9) + 7
#DIV = 998244353

def mod_inv_prime(a, mod=DIV):
    return pow(a, mod-2, mod)


def mod_inv(a, b):
    r = a
    w = b
    u = 1
    v = 0
    while w != 0:
        t = r//w
        r -= t*w
        r, w = w, r
        u -= t*v
        u, v = v, u
    return (u % b+b) % b


def CONV_TBL(max, mod=DIV):
    fac, finv, inv = [0]*max, [0]*max, [0]*max
    fac[0] = fac[1] = 1
    finv[0] = finv[1] = 1
    inv[1] = 1
    for i in range(2, max):
        fac[i] = fac[i-1]*i % mod
        inv[i] = mod - inv[mod % i] * (mod//i) % mod
        finv[i] = finv[i-1]*inv[i] % mod

    class CONV:
        def __init__(self):
            self.fac = fac
            self.finv = finv
            pass

        def ncr(self, n, k):
            if(n < k):
                return 0
            if(n < 0 or k < 0):
                return 0
            return fac[n]*(finv[k]*finv[n-k] % DIV) % DIV

    return CONV()


def cumsum(As):
    s = 0
    for a in As:
        s += a
        yield s


sys.setrecursionlimit(200005)


def dijkstra(G,start=0):
    heap = []
    cost = [INF]*len(G)
    heappush(heap,(0,start))
    while len(heap)!=0:
        c,n = heappop(heap)
        if cost[n] !=INF:
            continue
        cost[n]=c
        for e in G[n]:
            ec,v = e
            if cost[v]!=INF:
                continue
            heappush(heap,(ec+c,v))
    return cost


def main():

    N,K = isRaw()
    Ps = isRaw()
    Cs = isRaw()
    ans = max(Cs)
    for n in range(N):
        curPos = n
        cycles =[]
        visited = {}
        lastCost = 0
        vidx = 0
        loopStart = -1
        while True:
            visited[curPos] = vidx
            lastCost = lastCost+Cs[curPos]
            cycles.append(lastCost)
            curPos= Ps[curPos]-1
            
            if curPos in visited:
                loopStart = visited[curPos]
                break
        
        c_len = len(cycles)
        c_score = cycles[-1]
        if K>=c_len:
            ans = max(ans,c_score*(K//c_len))
        for c in range(0,min(c_len,K)):
            n_cycle = (K-c-1)//c_len
            score = cycles[c]+max(0,c_score*n_cycle)
            ans = max(score,ans)
    print (ans)

         
main()

Buggy submission (verdict: Runtime Error):
import bisect
import sys
import math
from collections import deque
from heapq import heappush, heappop

def sRaw():
    return input().rstrip("\r")


def iRaw():
    return int(input())


def ssRaw():
    return input().split()


def isRaw():
    return list(map(int, ssRaw()))


INF = 1 << 29

DIV = (10**9) + 7
#DIV = 998244353

def mod_inv_prime(a, mod=DIV):
    return pow(a, mod-2, mod)


def mod_inv(a, b):
    r = a
    w = b
    u = 1
    v = 0
    while w != 0:
        t = r//w
        r -= t*w
        r, w = w, r
        u -= t*v
        u, v = v, u
    return (u % b+b) % b


def CONV_TBL(max, mod=DIV):
    fac, finv, inv = [0]*max, [0]*max, [0]*max
    fac[0] = fac[1] = 1
    finv[0] = finv[1] = 1
    inv[1] = 1
    for i in range(2, max):
        fac[i] = fac[i-1]*i % mod
        inv[i] = mod - inv[mod % i] * (mod//i) % mod
        finv[i] = finv[i-1]*inv[i] % mod

    class CONV:
        def __init__(self):
            self.fac = fac
            self.finv = finv
            pass

        def ncr(self, n, k):
            if(n < k):
                return 0
            if(n < 0 or k < 0):
                return 0
            return fac[n]*(finv[k]*finv[n-k] % DIV) % DIV

    return CONV()


def cumsum(As):
    s = 0
    for a in As:
        s += a
        yield s


sys.setrecursionlimit(200005)


def dijkstra(G,start=0):
    heap = []
    cost = [INF]*len(G)
    heappush(heap,(0,start))
    while len(heap)!=0:
        c,n = heappop(heap)
        if cost[n] !=INF:
            continue
        cost[n]=c
        for e in G[n]:
            ec,v = e
            if cost[v]!=INF:
                continue
            heappush(heap,(ec+c,v))
    return cost


def main():
    N,K = isRaw()
    Ps = isRaw()
    Cs = isRaw()
    ans = -INF
    for n in range(N):
        curPos = n
        cycles =[]
        visited = {}
        lastCost = 0
        vidx = 0
        loopStart = -1
        while True:
            visited[curPos] = vidx
            cycles.append(lastCost+Cs[curPos])
            curPos= Ps[curPos]-1
            lastCost = cycles[-1]
            if curPos in visited:
                loopStart = visited[curPos]
                break
        lc = len(cycles)
        divK = K//lc
        score = 0
        print(n,cycles)
        if cycles[-1]>0:
            score+=cycles[-1]*divK
        score += max(cycles[:K % lc] if K % lc != 0 else cycles)

        ans = max(score,ans)
    print (ans)

         
main()