Problem p02585 (python)

Fixed submission:
# 公式解説を見てからやった
# 複数サイクルでき、1-K回移動可能
# N < 5000?多分、O(N * 2N)くらいで行けそうに見えるが。。。

if __name__ == "__main__":

	N,K = map(int,input().split())
	P = [ int(p)-1 for p in input().split() ]
	C = list(map(int,input().split()))

	# 一度計算したサイクル情報を一応キャッシュしておく。。。
	# あんまり意味なさそう
	cycleIDs = [ -1 for _ in range(N) ]
	cycleInfs = []


	# print(P)

	ans = -1e10
	cycleID = 0

	for n in range(N):

		v = n
		currentCycleItemCnt = 0
		currentCycleTotal = 0

		if cycleIDs[v] != -1:
			currentCycleItemCnt, currentCycleTotal = cycleInfs[ cycleIDs[v] ]
			# print(currentCycleItemCnt, currentCycleTotal)
		else:
			while True:
				# 全頂点について、属するサイクルを計算する
				currentCycleItemCnt += 1
				currentCycleTotal += C[v]

				v = P[v]
				if v == n:
					# サイクル発見
					cycleInfs.append( (currentCycleItemCnt, currentCycleTotal) )
					cycleID += 1
					break

				# 一応、一度サイクルを計算した頂点については、
				# その頂点の属するサイクルの情報をメモっておく。。。
				cycleIDs[v] = cycleID

		procCnt = 0
		currentCycleSumTmp = 0

		while True:

			# 頂点vにコマが置かれた時の最高スコアを計算し、
			# これまでの最高スコアを上回ったら、これまでの最高スコアを更新する
			procCnt += 1
			currentCycleSumTmp += C[v]
			if K < procCnt:
				break

			cycleLoopCnt = 0
			if procCnt < K and 0 < currentCycleTotal:
				cycleLoopCnt = ( K - procCnt ) // currentCycleItemCnt

			# print("v=", v, "currentCycleSumTmp=", currentCycleSumTmp, procCnt)
			tmp = currentCycleSumTmp + cycleLoopCnt * currentCycleTotal
			ans = max( ans, tmp )

			v = P[v]
			if v == n:
				# サイクル終了
				break

	print(ans)

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

if __name__ == "__main__":

	N,K = map(int,input().split())
	P = [ int(p)-1 for p in input().split() ]
	C = list(map(int,input().split()))

	# print(P)
	# 一度計算したサイクル情報をキャッシュしておくための配列
	cycleIDs = np.full( N, -1, np.int64 )
	cycleInfs = []
	cycleID = 0
	procCnt = 0

	for n in range(N):
		v = n

		if cycleIDs[v] != -1:
			continue
		else:
			currentCycleCosts = []
			while True:
				# 全頂点について、属するサイクルを計算する
				currentCycleCosts.append( C[v] )
				cycleIDs[v] = cycleID

				v = P[v]
				if cycleIDs[v] != -1:
					# サイクル終了
					# ループを含めない最大の処理回数
					procCnt = K % len( currentCycleCosts )
					# ただし、割り切れる場合も、1サイクル分は処理するものとする
					# (1回以上はやらなければならないのを考慮？)
					if procCnt == 0:
						procCnt = len( currentCycleCosts )

					cycleInfs.append( ( procCnt, np.array( currentCycleCosts ) ) )
					cycleID += 1
					break



	scores = []
	# loopScore = 0
	# procCnt = 0
	for procCnt, currentCycleCosts in cycleInfs:

		# サイクル内でループしてスコアを稼ぐ場合の考慮
		# if np.sum(currentCycleCosts) > 0:
		# 	cycleLoopCnt =  ( K - procCnt ) // currentCycleCosts.size
		# 	loopScore = cycleLoopCnt * np.sum(currentCycleCosts)

		# このサイクルに属する全頂点分をまとめて計算する
		for i in range(currentCycleCosts.size):
			if np.sum(currentCycleCosts) > 0:
				# acc = list(itertools.accumulate( np.roll( currentCycleCosts, i )[:procCnt] ) )
				scores.append( np.roll( currentCycleCosts, i )[:procCnt].cumsum().max() ) + (( K - procCnt ) // currentCycleCosts.size) * np.sum(currentCycleCosts) )
			else:
				scores.append( np.roll( currentCycleCosts, i )[:procCnt].cumsum().max() )


	print(max(scores))

Buggy submission (verdict: Runtime Error):
# 公式解説を見てからやった
# 複数サイクルでき、1-K回移動可能
# N < 5000?多分、O(N * 2N)くらいで行けそうに見えるが。。。

if __name__ == "__main__":

	N,K = map(int,input().split())
	P = [ int(p)-1 for p in input().split() ]
	C = list(map(int,input().split()))

	# 一度計算したサイクル情報を一応キャッシュしておく。。。
	# あんまり意味なさそう
	cycleIDs = [ -1 for _ in range(N) ]
	cycleInfs = []


	# print(P)

	ans = -1e19
	cycleID = 0

	for n in range(N):

		v = n
		currentCycleItemCnt = 0
		currentCycleTotal = 0

		if cycleIDs[v] != -1:
			currentCycleItemCnt, currentCycleTotal = cycleInfs[ cycleIDs[v] ]
			# print(currentCycleItemCnt, currentCycleTotal)
		else:
			while True:
				# 全頂点について、属するサイクルを計算する
				currentCycleItemCnt += 1
				currentCycleTotal += C[v]

				v = P[v]
				if v == n:
					# サイクル発見
					cycleInfs.append( (currentCycleItemCnt, currentCycleTotal) )
					cycleID += 1
					break

				# 一応、一度サイクルを計算した頂点については、
				# その頂点の属するサイクルの情報をメモっておく。。。
				cycleIDs[v] = cycleID

		procCnt = 0
		currentCycleSumTmp = 0

		while True:

			# 頂点vにコマが置かれた時の最高スコアを計算し、
			# これまでの最高スコアを上回ったら、これまでの最高スコアを更新する
			procCnt += 1
			currentCycleSumTmp += C[v]
			if K < procCnt:
				break

			cycleLoopCnt = 0
			if procCnt < K and 0 < currentCycleTotal:
				cycleLoopCnt = ( K - procCnt ) // currentCycleItemCnt

			print("v=", v, "currentCycleSumTmp=", currentCycleSumTmp, "procCnt, cycleLoopCnt, currentCycleTotal=", procCnt, cycleLoopCnt, currentCycleTotal)
			tmp = currentCycleSumTmp + cycleLoopCnt * currentCycleTotal
			ans = max( ans, tmp )

			v = P[v]
			if v == n:
				# サイクル終了
				break

	print(ans)