Problem p02585 (python)

Fixed submission:
import itertools
N, K = [int(_) for _ in input().split()]
P = [int(_) - 1 for _ in input().split()]
C = [int(_) for _ in input().split()]


class SegmentTree():
    def __init__(self, array, f, ti):
        """
        Parameters
        ----------
        array : list
            to construct segment tree from
        f : func
            binary operation of the monoid
        ti : 
            identity element of the monoid
        """
        self.f = f
        self.ti = ti
        self.n = n = 2**(len(array).bit_length())
        self.dat = dat = [ti] * n + array + [ti] * (n - len(array))
        for i in range(n - 1, 0, -1):  # build
            dat[i] = f(dat[i << 1], dat[i << 1 | 1])

    def update(self, p, v):  # set value at position p (0-indexed)
        f, n, dat = self.f, self.n, self.dat
        p += n
        dat[p] = v
        while p > 1:
            p >>= 1
            dat[p] = f(dat[p << 1], dat[p << 1 | 1])

    def operate_right(self, p, v):  # apply operator from the right side
        f, n, dat = self.f, self.n, self.dat
        p += n
        dat[p] = f(dat[p], v)
        while p > 1:
            p >>= 1
            dat[p] = f(dat[p << 1], dat[p << 1 | 1])

    def operate_left(self, p, v):  # apply operator from the left side
        f, n, dat = self.f, self.n, self.dat
        p += n
        dat[p] = f(v, dat[p])
        while p > 1:
            p >>= 1
            dat[p] = f(dat[p << 1], dat[p << 1 | 1])

    def query(self, l, r):  # result on interval [l, r) (0-indexed)
        f, ti, n, dat = self.f, self.ti, self.n, self.dat
        vl = vr = ti
        l += n
        r += n
        while l < r:
            if l & 1:
                vl = f(vl, dat[l])
                l += 1
            if r & 1:
                r -= 1
                vr = f(dat[r], vr)
            l >>= 1
            r >>= 1
        return f(vl, vr)


ans = -10**20
se = set()
for start in range(N):
    if start in se:
        continue
    now = start
    scores = []
    while True:
        if now in se:
            break
        se.add(now)
        scores += [C[now]]
        now = P[now]
    cyclelen = len(scores)
    cum = list(itertools.accumulate(scores * 2))
    st = SegmentTree(cum, max, -10**20)
    cycle, residue = divmod(K, cyclelen)
    ans = max([ans] + [
        st.query(start2 + 1, start2 + 1 + min(cyclelen, K)) - cum[start2]
        for start2 in range(cyclelen)
    ])
    if cum[-1] > 0 and cycle > 0:
        v1 = cum[-1] * (cycle - 1) // 2 + max(
            st.query(start2, start2 + cyclelen + 1) - cum[start2]
            for start2 in range(cyclelen))
        v2 = cum[-1] * cycle // 2 + max(
            st.query(start2, start2 + residue + 1) - cum[start2]
            for start2 in range(cyclelen))
        ans = max([ans, v1, v2])
print(ans)

Buggy submission (verdict: Runtime Error):
import itertools
N, K = [int(_) for _ in input().split()]
P = [int(_) - 1 for _ in input().split()]
C = [int(_) for _ in input().split()]


class SegmentTree():
    def __init__(self, array, f, ti):
        """
        Parameters
        ----------
        array : list
            to construct segment tree from
        f : func
            binary operation of the monoid
        ti : 
            identity element of the monoid
        """
        self.f = f
        self.ti = ti
        self.n = n = 2**(len(array).bit_length())
        self.dat = dat = [ti] * n + array + [ti] * (n - len(array))
        for i in range(n - 1, 0, -1):  # build
            dat[i] = f(dat[i << 1], dat[i << 1 | 1])

    def update(self, p, v):  # set value at position p (0-indexed)
        f, n, dat = self.f, self.n, self.dat
        p += n
        dat[p] = v
        while p > 1:
            p >>= 1
            dat[p] = f(dat[p << 1], dat[p << 1 | 1])

    def operate_right(self, p, v):  # apply operator from the right side
        f, n, dat = self.f, self.n, self.dat
        p += n
        dat[p] = f(dat[p], v)
        while p > 1:
            p >>= 1
            dat[p] = f(dat[p << 1], dat[p << 1 | 1])

    def operate_left(self, p, v):  # apply operator from the left side
        f, n, dat = self.f, self.n, self.dat
        p += n
        dat[p] = f(v, dat[p])
        while p > 1:
            p >>= 1
            dat[p] = f(dat[p << 1], dat[p << 1 | 1])

    def query(self, l, r):  # result on interval [l, r) (0-indexed)
        f, ti, n, dat = self.f, self.ti, self.n, self.dat
        vl = vr = ti
        l += n
        r += n
        while l < r:
            if l & 1:
                vl = f(vl, dat[l])
                l += 1
            if r & 1:
                r -= 1
                vr = f(dat[r], vr)
            l >>= 1
            r >>= 1
        return f(vl, vr)


ans = -10**20
se = set()
for start in range(N):
    if start in se:
        continue
    now = start
    scores = []
    while True:
        if now in se:
            break
        se.add(now)
        scores += [C[now]]
        now = P[now]
    cyclelen = len(scores)
    cycle, residue = divmod(K, cyclelen)
    scores *= 2
    scores = list(itertools.accumulate(scores))
    st = SegmentTree(scores, max, -10**20)
    ans = max(
        [ans] +
        [st.query(start2 + 1, start2 + K + 1) for start2 in range(cyclelen)])
    if scores[-1] > 0 and cycle > 0:
        ans2 = scores[-1] * cycle // 2
        v = max(
            st.query(start2 + 1, start2 + residue + 1) - scores[start2]
            for start2 in range(cyclelen))
        ans = max([ans, ans2, ans2 + v])

print(ans)

Buggy submission (verdict: Runtime Error):
import itertools
N, K = [int(_) for _ in input().split()]
P = [int(_) - 1 for _ in input().split()]
C = [int(_) for _ in input().split()]


class SegmentTree():
    def __init__(self, array, f, ti):
        """
        Parameters
        ----------
        array : list
            to construct segment tree from
        f : func
            binary operation of the monoid
        ti : 
            identity element of the monoid
        """
        self.f = f
        self.ti = ti
        self.n = n = 2**(len(array).bit_length())
        self.dat = dat = [ti] * n + array + [ti] * (n - len(array))
        for i in range(n - 1, 0, -1):  # build
            dat[i] = f(dat[i << 1], dat[i << 1 | 1])

    def update(self, p, v):  # set value at position p (0-indexed)
        f, n, dat = self.f, self.n, self.dat
        p += n
        dat[p] = v
        while p > 1:
            p >>= 1
            dat[p] = f(dat[p << 1], dat[p << 1 | 1])

    def operate_right(self, p, v):  # apply operator from the right side
        f, n, dat = self.f, self.n, self.dat
        p += n
        dat[p] = f(dat[p], v)
        while p > 1:
            p >>= 1
            dat[p] = f(dat[p << 1], dat[p << 1 | 1])

    def operate_left(self, p, v):  # apply operator from the left side
        f, n, dat = self.f, self.n, self.dat
        p += n
        dat[p] = f(v, dat[p])
        while p > 1:
            p >>= 1
            dat[p] = f(dat[p << 1], dat[p << 1 | 1])

    def query(self, l, r):  # result on interval [l, r) (0-indexed)
        f, ti, n, dat = self.f, self.ti, self.n, self.dat
        vl = vr = ti
        l += n
        r += n
        while l < r:
            if l & 1:
                vl = f(vl, dat[l])
                l += 1
            if r & 1:
                r -= 1
                vr = f(dat[r], vr)
            l >>= 1
            r >>= 1
        return f(vl, vr)


ans = -float('inf')
se = set()
for start in range(N):
    if start in se:
        continue
    now = start
    scores = []
    while True:
        if now == start:
            break
        se.add(now)
        scores += [C[now]]
        now = P[now]
    cyclelen = len(scores)
    cycle, residue = divmod(K, cyclelen)
    scores *= 3
    scores = list(itertools.accumulate(scores))
    st = SegmentTree(scores, max, -float('inf'))
    ans = max([ans] + [
        st.query(start2 + 1, start2 + min(K + 1, cyclelen)) - scores[start2]
        for start2 in range(cyclelen)
    ])
    scorelastcycle = scores[-1] * (cycle - 1) // 3
    if scorelastcycle > 0 and cycle > 0:
        d = 0
        for start2 in range(cyclelen):
            d = max(
                d,
                st.query(start2, start2 + cyclelen + residue + 1) -
                scores[start2])
        scorelastcycle += d
        ans = max(ans, scorelastcycle)
print(ans)

Buggy submission (verdict: Runtime Error):
import itertools
N, K = [int(_) for _ in input().split()]
P = [int(_) - 1 for _ in input().split()]
C = [int(_) for _ in input().split()]


class SegmentTree():
    def __init__(self, array, f, ti):
        """
        Parameters
        ----------
        array : list
            to construct segment tree from
        f : func
            binary operation of the monoid
        ti : 
            identity element of the monoid
        """
        self.f = f
        self.ti = ti
        self.n = n = 2**(len(array).bit_length())
        self.dat = dat = [ti] * n + array + [ti] * (n - len(array))
        for i in range(n - 1, 0, -1):  # build
            dat[i] = f(dat[i << 1], dat[i << 1 | 1])

    def update(self, p, v):  # set value at position p (0-indexed)
        f, n, dat = self.f, self.n, self.dat
        p += n
        dat[p] = v
        while p > 1:
            p >>= 1
            dat[p] = f(dat[p << 1], dat[p << 1 | 1])

    def operate_right(self, p, v):  # apply operator from the right side
        f, n, dat = self.f, self.n, self.dat
        p += n
        dat[p] = f(dat[p], v)
        while p > 1:
            p >>= 1
            dat[p] = f(dat[p << 1], dat[p << 1 | 1])

    def operate_left(self, p, v):  # apply operator from the left side
        f, n, dat = self.f, self.n, self.dat
        p += n
        dat[p] = f(v, dat[p])
        while p > 1:
            p >>= 1
            dat[p] = f(dat[p << 1], dat[p << 1 | 1])

    def query(self, l, r):  # result on interval [l, r) (0-indexed)
        f, ti, n, dat = self.f, self.ti, self.n, self.dat
        vl = vr = ti
        l += n
        r += n
        while l < r:
            if l & 1:
                vl = f(vl, dat[l])
                l += 1
            if r & 1:
                r -= 1
                vr = f(dat[r], vr)
            l >>= 1
            r >>= 1
        return f(vl, vr)


ans = -float('inf')
se = set()
for start in range(N):
    if start in se:
        continue
    now = start
    scores = []
    while True:
        if now in se:
            break
        se.add(now)
        scores += [C[now]]
        now = P[now]
    cyclelen = len(scores)
    cycle, residue = divmod(K, cyclelen)
    scores *= 2
    scores = list(itertools.accumulate(scores))
    st = SegmentTree(scores, max, -float('inf'))
    ans = max([ans] + [
        st.query(start2, start2 + K + 1) - scores[start2]
        for start2 in range(cyclelen)
    ])
    scorelastcycle = scores[-1] * cycle // 2
    if scorelastcycle > 0 and cycle > 0:
        d = 0
        for start2 in range(cyclelen):
            d = max(d, st.query(start2, start2 + residue + 1) - scores[start2])
        scorelastcycle += d
        ans = max(ans, scorelastcycle)
print(ans)

Buggy submission (verdict: Runtime Error):
N, K = [int(_) for _ in input().split()]
P = [int(_) - 1 for _ in input().split()]
C = [int(_) for _ in input().split()]
child_parent = [0] * N
ans = -float('inf')
for p, c in enumerate(P):
    child_parent[c] = p
for i in range(N):
    f = i
    i_cnt = [None] * N
    i_cnt[i] = 0
    cnt_score = [0]
    score = 0
    cnt = 0
    while True:
        i = P[i]
        cnt += 1
        score += C[i]
        cnt_score += [score]
        if cnt > K or i_cnt[i] is not None:
            break
        i_cnt[i] = cnt
    score1 = cnt_score[i_cnt[i]]
    cyclelen = cnt - i_cnt[i]
    if cyclelen:
        q, r = divmod((K - i_cnt[i]), cyclelen)
        ans1 = score1 + (score -
                         score1) * q + max([0] + cnt_score[i_cnt[i]:][:r + 1])
        ans2 = max([C[f]] + cnt_score[:i_cnt[i]])
        ans = max([ans, ans1, ans2])
    else:
        ans = max([ans] + cnt_score)
print(ans)