Problem p02618 (python)

Fixed submission:
inf = 10**12


def main():
    import sys
    from bisect import bisect_left
    input = sys.stdin.readline

    D = int(input())
    C = list(map(int, input().split()))
    S = []
    for _ in range(D):
        S.append(list(map(int, input().split())))
    ans = 0
    last = [-1] * 26
    days = [[-1] for _ in range(26)]
    T = []
    # greedy
    for d in range(D):
        plus = -inf
        t = -1
        for tt in range(26):
            plus_new = S[d][tt] + C[tt] * (d - last[tt])
            if plus_new > plus:
                t = tt
                plus = plus_new
        T.append(t)
        ans += S[d][t]
        last[t] = d
        days[t].append(d)
        for i in range(26):
            ans -= C[i] * (d - last[i])
    for t in range(26):
        days[t].append(D)
    for t in T:
        print(t+1)
    exit()







    M = int(input())
    for _ in range(M):
        d, q = map(int, input().split())
        d -= 1
        q -= 1
        t = T[d]
        ans += S[d][q]
        ans -= S[d][t]
        i = bisect_left(days[q], d)
        ans += C[q] * (d - days[q][i-1]) * (days[q][i] - d)
        j = bisect_left(days[t], d)
        ans -= C[t] * (d - days[t][j-1]) * (days[t][j+1] - d)
        #print(days[q], days[t])
        T[d] = q
        days[q].append(d)
        days[q].sort()
        days[t].remove(d)
        print(ans)


if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
inf = 10**12


def main():
    import sys
    from random import randint
    from bisect import bisect_left
    input = sys.stdin.readline

    D = int(input())
    C = list(map(int, input().split()))
    S = []
    for _ in range(D):
        S.append(list(map(int, input().split())))
    ans = 0
    last = [-1] * 26
    days = [[-1] for _ in range(26)]
    T = []
    # greedy
    for d in range(D):
        plus = -inf
        t = -1
        for tt in range(26):
            plus_new = S[d][tt] + C[tt] * (d - last[tt])
            if plus_new > plus:
                t = tt
                plus = plus_new
        T.append(t)
        ans += S[d][t]
        last[t] = d
        days[t].append(d)
        for i in range(26):
            ans -= C[i] * (d - last[i])
    for t in range(26):
        days[t].append(D)

    def calc(d, q, cur):
        t = T[d]
        cur += S[d][q]
        cur -= S[d][t]
        i = bisect_left(days[q], d)
        cur += C[q] * (d - days[q][i - 1]) * (days[q][i] - d)
        j = bisect_left(days[t], d)
        cur -= C[t] * (d - days[t][j - 1]) * (days[t][j + 1] - d)
        return cur

    # climb
    size = 5
    try_num = 1000
    for x in range(D // size):
        d_list = list(range(x*size, (x+1)*size))
        for _ in range(try_num):
            t_new = [randint(0, 25) for _ in range(size)]
            ans_new = ans
            for d in d_list:
                ans_new = calc(d, t_new[d], ans_new)
            if ans_new > ans:
                ans = ans_new
                for d in d_list:
                    t_old = T[d]
                    T[d] = t_new[d]
                    days[t_new[d]].append(d)
                    days[t_new[d]].sort()
                    days[t_old].remove(d)
    for t in T:
        print(t+1)


if __name__ == '__main__':
    main()