Problem p02585 (python)

Fixed submission:
N, K = map(int, input().split())
P = [0] + list(map(int, input().split()))
C = [0] + list(map(int, input().split()))

# Pが1,2,...Nの順列なので、1,2,3,4,5,3,4,5,...のような途中からのループは起こり得ない
# スタートの点数が入らないのは面倒なので、スタートの次のマスから考える

ans = -(10**9)
# どのマスから始めるか
for i in range(1, N + 1):
    # ループに入るまでの経路(iの次から)
    route = []
    visit = [False] * (N + 1)
    visit[i] = True
    now = i
    while True:
        next = P[now]
        # まだ訪れていないマスの場合
        if visit[next] == False:
            visit[next] = True
            route.append(next)
            now = next
        # 既に訪れたマスの場合
        else:
            break
    # iに戻ってくる
    route.append(i)

    cumsum = [0]
    for x in route:
        cumsum.append(cumsum[-1] + C[x])
    # print(cumsum)

    # ループするほど損する場合
    if cumsum[-1] <= 0:
        if K >= len(route):
            tmp = max(cumsum[1:])
        else:
            tmp = max(cumsum[1:K + 1])
        ans = max(ans, tmp)
    # ループするほど得する場合
    else:
        loop_size = len(route)
        loop_cnt = K // loop_size
        spare_cnt = K % loop_size
        # (loop_cnt -1)回ループしたところから
        cumsum2 = [cumsum[-1] * (loop_cnt - 1)]
        for x in route:
            cumsum2.append(cumsum2[-1] + C[x])
        for x in route[:spare_cnt]:
            cumsum2.append(cumsum2[-1] + C[x])
        tmp = max(cumsum2[1:])
        ans = max(ans, tmp)

print(ans)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
P = [0] + list(map(int, input().split()))
C = [0] + list(map(int, input().split()))

# Pが1,2,...Nの順列なので、1,2,3,4,5,3,4,5,...のような途中からのループは起こり得ない
# スタートの点数が入らないのは面倒なので、スタートの次のマスから考える

ans = -(10**9)
# どのマスから始めるか
for i in range(1, N + 1):
    # ループに入るまでの経路(iの次から)
    route = []
    visit = [False] * (N + 1)
    visit[i] = True
    now = i
    while True:
        next = P[now]
        # まだ訪れていないマスの場合
        if visit[next] == False:
            visit[next] = True
            route.append(next)
            now = next
        # 既に訪れたマスの場合
        else:
            break
    # iに戻ってくる
    route.append(i)

    cumsum = [0]
    for x in route:
        cumsum.append(cumsum[-1] + C[x])
    # print(cumsum)

    # ループするほど損する場合
    if cumsum[-1] <= 0:
        tmp = max(cumsum[1:])
        ans = max(ans, tmp)
    # ループするほど得する場合
    else:
        loop_size = len(route)
        loop_cnt = K // loop_size
        spare_cnt = K % loop_size
        tmp = cumsum[-1] * loop_cnt + max(cumsum[1:spare_cnt + 1])
        ans = max(ans,tmp)

print(ans)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
P = list(map(int, input().split()))
C = list(map(int, input().split()))

roots = []

# どのマスから始めるか
for i in range(1, N + 1):
    root = [i]
    visit = [False] * N
    visit[i - 1] = True
    now = i
    for j in range(K):
        next = P[now - 1]
        if visit[next - 1] == True:
            root.append(next)
            break
        else:
            visit[next - 1] = True
            root.append(next)
            now = next
    roots.append(root)
print(roots)

ans = 0
for root in roots:
    head = root[-1]
    head_idx = root.index(head)
    tail = root[-2]
    tail_idx = root.index(tail)
    roop_l = tail_idx - head_idx + 1
    roop = []
    for i in range(head_idx, tail_idx + 1):
        roop.append(root[i])
    # print(roop,"roop")
    if roop == []:
        cumsum = [0]
        for i in range(1, len(root)):
            cumsum.append(cumsum[-1] + C[root[i] - 1])
        # print(cumsum,"cumsum")
        m = max(cumsum)
        # print(m)
        ans = max(ans, m)
    else:
        cumsum = [0]
        for item in roop:
            cumsum.append(cumsum[-1]+C[item-1])
        m = max(cumsum)
        max_idx = cumsum.index(m)
        # print(cumsum,"else",max_idx)
        #
        cumsum2 = [0]
        for i in range(1,len(root)-1):
            cumsum2.append(cumsum2[-1]+C[root[i]-1])
        # print(cumsum2)
        if cumsum[-1]<=0:
            for item in root:
                cumsum2.append(cumsum[-1]+C[item-1])
            ans = max(cumsum2)
        else:
            K -= len(cumsum2)
            roop_cnt = K//len(roop)
            cumsum2.append(cumsum2[-1]+ cumsum[-1]*roop_cnt + m)
            ans = max(cumsum2)
print(ans)