Problem p02585 (python)

Fixed submission:
import numpy as np
 
n, k = map(int, input().split())
p = list(map(int, input().split()))
c = list(map(int, input().split()))
 
cycles, check = [], [False]*n
for pi in p:
    if not check[pi-1]:
        curr, cycle = pi, []
        check[pi-1] = True
        while pi != p[curr-1]:
            check[curr-1] = True
            cycle.append(c[curr-1])
            curr = p[curr-1]
        check[curr-1] = True
        cycles.append(cycle + [c[curr-1]])
 
scores = []
for cycle in (np.array(cycle) for cycle in cycles):
    for i in range(cycle.size):
        cum = np.roll(cycle, i).cumsum()
        if cum[-1] > 0 and k > cycle.size:
            scores.append(cum[-1]*(k//cycle.size-1) + np.concatenate([cum, cum[:k % cycle.size]+cum[-1]]).max())
        else:
            scores.append(cum[:min(k, cum.size)].max())
print(max(scores))

Buggy submission (verdict: Runtime Error):
import numpy as np

n, k = map(int, input().split())
p = list(map(int, input().split()))
c = list(map(int, input().split()))

cycles, check = [], [False]*n
for pi in p:
    if not check[pi-1]:
        curr, cycle = pi, []
        check[pi-1] = True
        while pi != p[curr-1]:
            check[curr-1] = True
            cycle.append(c[curr-1])
            curr = p[curr-1]
        check[curr-1] = True
        cycles.append(cycle + [c[curr-1]])

scores = []
for cycle in (np.array(cycle) for cycle in cycles):
    for i in range(cycle.size):
        cum = np.roll(cycle, i).cumsum()
        if cum[-1] > 0 and k > cycle.size:
            scores.append(cum[-1]*(k//cycle.size-1) + np.concatenate([cum, cum[:k % cycle.size]+cum[-1]]).max())
        else:
            scores.append(cum[:min(k, cum.size)-1].max())
print(max(scores))

Buggy submission (verdict: Runtime Error):
n, k = map(int, input().split())
p = list(map(int, input().split()))
c = list(map(int, input().split()))

maxcost = -float('inf')
for pi in p:    # O(n^2)
    curr, costs = pi, [0]
    while pi != p[curr-1]:
        costs.append(costs[-1]+c[curr-1])
        curr = p[curr-1]
    costs.append(costs[-1]+c[curr-1])
    l = len(costs)-1
    if costs[-1] > 0 and k > l:
        maxcost = max(maxcost, costs[-1]*(k//l)+max(costs[1:k % l+1]) or 0)
    else:
        maxcost = max(maxcost, max(costs[1:min(k, l)+1]))

print(maxcost)

Buggy submission (verdict: Runtime Error):
n, k = map(int, input().split())
p = list(map(int, input().split()))
c = list(map(int, input().split()))

maxcost = -float('inf')
for pi in p:    # O(n^2)
    curr, costs = pi, [0]
    while pi != p[curr-1]:
        costs.append(costs[-1]+c[curr-1])
        curr = p[curr-1]
    costs.append(costs[-1]+c[curr-1])
    l = len(costs)-1
    if costs[-1] > 0:
        maxcost = max(maxcost, costs[-1]*(k//l)+max(costs[1:k % l+1]))
    else:
        maxcost = max(maxcost, max(costs[1:min(k, l)+1]))

print(maxcost)

Buggy submission (verdict: Runtime Error):
n, k = map(int, input().split())
p = list(map(int, input().split()))
c = list(map(int, input().split()))

maxcost = -float('inf')
for pi in p:    # O(n^2)
    curr, costs = pi, [0]
    while pi != p[curr-1]:
        curr = p[curr-1]
        costs.append(costs[-1]+c[curr-1])
    loopsum = costs[-1]+c[pi-1]
    if loopsum > 0:
        maxcost = max(maxcost, loopsum*(k//len(costs))+max(costs[1:k % len(costs)+1]))
    else:
        maxcost = max(maxcost, max(costs[1:min(k, n)+1]))

print(maxcost)

Buggy submission (verdict: Runtime Error):
n, k = map(int, input().split())
p = list(map(int, input().split()))
c = list(map(int, input().split()))

maxcost = -float('inf')
for pi in p:    # O(n^2)
    curr, costs = pi, [0]
    while pi != p[curr-1]:
        curr = p[curr-1]
        costs.append(costs[-1]+c[curr-1])
    loopsum = costs[-1]+c[pi-1]
    if loopsum > 0:
        maxcost = max(maxcost, loopsum*(k//len(costs))+max(costs[1:k % len(costs)+1]))
    else:
        maxcost = max(maxcost, max(costs[1:min(k, n)]))

print(maxcost)