Problem p02585 (python)

Fixed submission:
import bisect, collections, copy, heapq, itertools, math, string, sys
input = lambda: sys.stdin.readline().rstrip() 
sys.setrecursionlimit(10**7)
INF = float('inf')
def I(): return int(input())
def F(): return float(input())
def SS(): return input()
def LI(): return [int(x) for x in input().split()]
def LI_(): return [int(x)-1 for x in input().split()]
def LF(): return [float(x) for x in input().split()]
def LSS(): return input().split()

def resolve():
    N, K = LI()
    P = LI_()
    C = LI()

    ans = -INF
    for i in range(N):
        # iを含むサイクルの探索
        cycle = []
        c = i
        while True:
            c = P[c]
            cycle.append(c)
            if c == i:
                break
        len_cycle = len(cycle)
        # print(i, cycle, [C[j] for j in cycle])

        # 最大スコアの探索
        score = 0
        cycle_score_total = sum([C[j] for j in cycle])
        cycle_score_cum = list(itertools.accumulate([C[j] for j in cycle]))
        # 終了地点をjと決める
        for j in range(min(len_cycle, K)):
            score = cycle_score_cum[j]
            if cycle_score_total > 0:
                score += (K - j - 1) // len_cycle * cycle_score_total
            # print(j, score)
            ans = max(score, ans)

    print(ans)

if __name__ == '__main__':
    resolve()

Buggy submission (verdict: Runtime Error):
import bisect, collections, copy, heapq, itertools, math, string, sys
input = lambda: sys.stdin.readline().rstrip() 
sys.setrecursionlimit(10**7)
INF = float('inf')
def I(): return int(input())
def F(): return float(input())
def SS(): return input()
def LI(): return [int(x) for x in input().split()]
def LI_(): return [int(x)-1 for x in input().split()]
def LF(): return [float(x) for x in input().split()]
def LSS(): return input().split()

def resolve():
    N, K = LI()
    P = LI_()
    C = LI()

    ans = -INF
    for i in range(N):
        # iを含むサイクルの探索
        cycle = []
        c = i
        while True:
            c = P[c]
            cycle.append(c)
            if c == i:
                break
        len_cycle = len(cycle)

        # 最大スコアの探索
        score = 0

        # Kの残り
        r = K
        cycle_score_total = sum([C[j] for j in cycle])
        if len_cycle < K:
            # 周回して増えるなら周回する
            if cycle_score_total > 0:
                score += K // len_cycle * cycle_score_total
                # 周回後の余り
                r = K % len_cycle
            # 周回して増えないなら周回しない
            else:
                # サイクル一周までの範囲でどこまでもいける
                r = len_cycle

        # Kの残っている範囲で最大スコア探索
        cycle_score = [C[j] for j in cycle[:r]]
        cycle_score_cum = list(itertools.accumulate(cycle_score))
        score += max(cycle_score_cum)

        ans = max(score, ans)

    print(ans)

if __name__ == '__main__':
    resolve()

Buggy submission (verdict: Runtime Error):
import bisect, collections, copy, heapq, itertools, math, string, sys
input = lambda: sys.stdin.readline().rstrip() 
sys.setrecursionlimit(10**7)
INF = float('inf')
def I(): return int(input())
def F(): return float(input())
def SS(): return input()
def LI(): return [int(x) for x in input().split()]
def LI_(): return [int(x)-1 for x in input().split()]
def LF(): return [float(x) for x in input().split()]
def LSS(): return input().split()

def resolve():
    N, K = LI()
    P = LI_()
    C = LI()

    ans = -INF
    for i in range(N):
        # iを含むサイクルの探索
        cycle = []
        c = i
        while True:
            c = P[c]
            cycle.append(c)
            if c == i:
                break
        len_cycle = len(cycle)

        # 最大スコアの探索
        score = 0

        # Kの残り
        r = K
        cycle_score_total = sum([C[j] for j in cycle])
        if len_cycle < K:
            # 周回して増えるなら周回する
            if cycle_score_total > 0:
                score += K // len_cycle * cycle_score_total
                # 周回後の余り
                r = K % len_cycle
            # 周回して増えないなら周回しない
            else:
                # サイクル一周までの範囲でどこまでもいける
                r = len_cycle

        # Kの残っている範囲で最大スコア探索
        cycle_score = [C[cycle[j]] for j in range(r)]
        # cycle_score_cum = list(itertools.accumulate(cycle_score))
        cycle_score_cum = [0] * r
        cycle_score_cum[0] = cycle_score[0]
        for i in range(r - 1):
            cycle_score_cum[i+1] = cycle_score[i+1] + cycle_score_cum[i]
        score += max(cycle_score_cum)

        ans = max(score, ans)

    print(ans)

if __name__ == '__main__':
    resolve()

Buggy submission (verdict: Runtime Error):
import bisect, collections, copy, heapq, itertools, math, string, sys
input = lambda: sys.stdin.readline().rstrip() 
sys.setrecursionlimit(10**7)
INF = float('inf')
def I(): return int(input())
def F(): return float(input())
def SS(): return input()
def LI(): return [int(x) for x in input().split()]
def LI_(): return [int(x)-1 for x in input().split()]
def LF(): return [float(x) for x in input().split()]
def LSS(): return input().split()

def resolve():
    N, K = LI()
    P = LI_()
    C = LI()

    ans = -INF
    for i in range(N):
        # iを含むサイクルの探索
        cycle = []
        c = i
        while True:
            c = P[c]
            cycle.append(c)
            if c == i:
                break
        len_cycle = len(cycle)

        # 最大スコアの探索
        score = 0

        # Kの残り
        r = K
        cycle_score_total = sum([C[j] for j in cycle])
        if len_cycle < K:
            # 周回して増えるなら周回する
            if cycle_score_total > 0:
                score += K // len_cycle * cycle_score_total
                # 周回後の余り
                r = K % len_cycle
            # 周回して増えないなら周回しない
            else:
                # サイクル一周までの範囲でどこまでもいける
                r = len_cycle

        # Kの残っている範囲で最大スコア探索
        cycle_score = [C[cycle[j]] for j in range(r)]
        cycle_score_cum = list(itertools.accumulate(cycle_score))
        score += max(cycle_score_cum)

        ans = max(score, ans)

    print(ans)

if __name__ == '__main__':
    resolve()

Buggy submission (verdict: Runtime Error):
import bisect, collections, copy, heapq, itertools, math, string, sys
input = lambda: sys.stdin.readline().rstrip() 
sys.setrecursionlimit(10**7)
INF = float('inf')
def I(): return int(input())
def F(): return float(input())
def SS(): return input()
def LI(): return [int(x) for x in input().split()]
def LI_(): return [int(x)-1 for x in input().split()]
def LF(): return [float(x) for x in input().split()]
def LSS(): return input().split()

def resolve():
    N, K = LI()
    P = LI_()
    C = LI()

    ans = -INF
    for i in range(N):
        # iを含むサイクルの探索
        cycle = []
        c = i
        while True:
            c = P[c]
            cycle.append(c)
            if c == i:
                break
        len_cycle = len(cycle)

        # 最大スコアの探索
        score = 0

        # Kの残り
        r = K
        cycle_score_total = sum([C[j] for j in cycle])
        if len_cycle < K:
            # 周回して増えるなら周回する
            if cycle_score_total > 0:
                score += K // len_cycle * cycle_score_total
                # 周回後の余り
                r = K % len_cycle
            # 周回して増えないなら周回しない
            else:
                # サイクル一周までの範囲でどこまでもいける
                r = len_cycle

        # Kの残っている範囲で最大スコア探索
        cycle_score = [C[cycle[j]] for j in range(r)]
        cycle_score_cum = list(itertools.accumulate(cycle_score))
        score += max(cycle_score_cum)

        ans = max(score, ans)

    print(ans)

if __name__ == '__main__':
    resolve()

Buggy submission (verdict: Runtime Error):
import bisect, collections, copy, heapq, itertools, math, string, sys
input = lambda: sys.stdin.readline().rstrip() 
sys.setrecursionlimit(10**7)
INF = float('inf')
def I(): return int(input())
def F(): return float(input())
def SS(): return input()
def LI(): return [int(x) for x in input().split()]
def LI_(): return [int(x)-1 for x in input().split()]
def LF(): return [float(x) for x in input().split()]
def LSS(): return input().split()

def resolve():
    N, K = LI()
    P = LI_()
    C = LI()

    # Pをサイクルに分割
    cycle_C = []
    visited = [False] * N
    for i in range(N):
        if not visited[i]:
            c = i
            tmp_C = [C[c]]
            while True:
                visited[c] = True
                c = P[c]
                if visited[c]:
                    cycle_C.append(tmp_C)
                    break
                tmp_C.append(C[c])
    # print(cycle_C)

    ans = -INF
    for i in cycle_C:
        cycle_sum = sum(i)
        # print(cycle_sum)

        # Kがサイクル長より大きく、1週の点数が正ならば、周回する
        if len(i) < K and cycle_sum > 0:
            # サイクルを2つくっつけて累積和を取り、その差を全探索
            i_cum = list(itertools.accumulate(i * 2, initial=0))
            tmp = -INF
            for j, k in itertools.combinations(range(len(i) * 2), 2):
                if k - j <= K % len(i):
                    tmp = max(i_cum[k] - i_cum[j], tmp)
            tmp += K // len(i) * cycle_sum 

        else:
            i_cum = list(itertools.accumulate(i * 2, initial=0))
            tmp = -INF
            for j, k in itertools.combinations(range(len(i) * 2), 2):
                if k - j <= K:
                    tmp = max(i_cum[k] - i_cum[j], tmp)

        ans = max(tmp, ans)

    print(ans)

if __name__ == '__main__':
    resolve()

Buggy submission (verdict: Runtime Error):
import bisect, collections, copy, heapq, itertools, math, string, sys
input = lambda: sys.stdin.readline().rstrip() 
sys.setrecursionlimit(10**7)
INF = float('inf')
def I(): return int(input())
def F(): return float(input())
def SS(): return input()
def LI(): return [int(x) for x in input().split()]
def LI_(): return [int(x)-1 for x in input().split()]
def LF(): return [float(x) for x in input().split()]
def LSS(): return input().split()

def resolve():
    N, K = LI()
    P = LI_()
    C = LI()

    # Pをサイクルに分割
    cycle_C = []
    visited = [False] * N
    for i in range(N):
        if not visited[i]:
            c = i
            tmp_C = [C[c]]
            while True:
                visited[c] = True
                c = P[c]
                if visited[c]:
                    cycle_C.append(tmp_C)
                    break
                tmp_C.append(C[c])
    # print(cycle_C)

    ans = -INF
    for i in cycle_C:
        cycle_sum = sum(i)
        # print(cycle_sum)

        # Kがサイクル長より大きく、1週の点数が正ならば、周回する
        if len(i) < K and cycle_sum > 0:
            # サイクルを2つくっつけて累積和を取り、その差を全探索
            i_cum = list(itertools.accumulate(i * 2, initial=0))
            tmp = -INF
            for j, k in itertools.combinations(range(len(i) * 2), 2):
                if k - j <= K % len(i):
                    tmp = max(i_cum[k] - i_cum[j], tmp)
            tmp += K // len(i) * cycle_sum 

        else:
            i_cum = list(itertools.accumulate(i * 2, initial=0))
            tmp = -INF
            for j, k in itertools.combinations(range(len(i) * 2), 2):
                if k - j <= K:
                    tmp = max(i_cum[k] - i_cum[j], tmp)

        ans = max(tmp, ans)

    print(ans)

if __name__ == '__main__':
    resolve()