Problem p03148 (python)

Fixed submission:
n, k = [int(i) for i in input().split(" ")]
t = []
d = []

for i in range(0, n):
    x, y = [int(i) for i in input().split(" ")]
    t.append(x)
    d.append(y)

dt = zip(d, t)
dt = sorted(dt, key = lambda k:k[0], reverse=True)
t_max = [False for i in range(0, n + 1)]
d_max = []
d_non_max = []

for i in range(0, n):
    if t_max[dt[i][1]]:
        d_non_max.append(dt[i][0])
    else:
        t_max[dt[i][1]] = True
        d_max.append(dt[i][0])
        
p = [0 for i in range(0, len(d_max) + 1)]
sum_max = 0
sum_non_max = sum(d_non_max[0:min(len(d_non_max), k)])

for i in range(1, len(p)):
    sum_max += d_max[i - 1]
    if k - i >= 0 and k - i < len(d_non_max):
        sum_non_max -= d_non_max[k - i]
    if k - i >= 0 and k - i - 1 < len(d_non_max):
        p[i] = sum_max + sum_non_max + i * i
    else:
        p[i] = 0
        
print(max(p))

Buggy submission (verdict: Runtime Error):
n, k = [int(i) for i in input().split(" ")]
t = []
d = []

for i in range(0, n):
    x, y = [int(i) for i in input().split(" ")]
    t.append(x)
    d.append(y)

dt = zip(d, t)
dt = sorted(dt, key = lambda k:k[0], reverse=True)
td = [[] for i in range(0, n + 1)]
tp = []
dp = [[0 for j in range(0, k)] for i in range(0, k)]
eaten = [[0 for j in range(0, k)] for i in range(0, k)]

for i in range(0, n):
    td[dt[i][1]].append(dt[i][0])
    if len(td[dt[i][1]]) == 1:
        tp.append(dt[i][1])

dp[0][0] = td[tp[0]][0] + 1
eaten[0][0] = 1

for j in range(1, k):
    if eaten[0][j - 1] < len(td[tp[0]]):
        dp[0][j] = dp[0][j - 1] + td[tp[0]][eaten[0][j - 1]]
        eaten[0][j] = eaten[0][j - 1] + 1
    else:
        break


for i in range(1, len(tp)):
    eaten[i][i] = 1
    dp[i][i] =  dp[i - 1][i - 1] + td[tp[i]][0] + 2 * i + 1
    for j in range(i + 1, k):
        if dp[i - 1][j - 1] == 0 or (eaten[i][j - 1] < len(td[tp[i]]) and 2 * i + 1 + td[tp[i]][0] < td[tp[i]][eaten[i][j - 1]]):
            if eaten[i][j - 1] < len(td[tp[i]]) and dp[i][j - 1] > 0:
                dp[i][j] = dp[i][j - 1] + td[tp[i]][eaten[i][j - 1]]
                eaten[i][j] = eaten[i][j - 1] + 1
        else:
            if dp[i - 1][j - 1] > 0:
                dp[i][j] = dp[i - 1][j - 1] + 2 * i + 1 + td[tp[i]][0]
                eaten[i][j] = 1

print(max([i[-1] for i in dp]))

Buggy submission (verdict: Runtime Error):
n, k = [int(i) for i in input().split(" ")]
t = []
d = []

for i in range(0, n):
    x, y = [int(i) for i in input().split(" ")]
    t.append(x)
    d.append(y)
    
dt = zip(d, t)
dt = sorted(dt, key = lambda k:k[0], reverse=True)
td = [[] for i in range(0, n + 1)]
tp = []
dp = [[0 for j in range(0, k)] for i in range(0, k)]
eaten = [[0 for j in range(0, k)] for i in range(0, k)]

for i in range(0, n):
    td[dt[i][1]].append(dt[i][0])
    if len(td[dt[i][1]]) == 1:
        tp.append(dt[i][1])

dp[0][0] = td[tp[0]][0] + 1
eaten[0][0] = 1

for j in range(1, k):
    if eaten[0][j - 1] < len(td[tp[0]]):
        dp[0][j] = dp[0][j - 1] + td[tp[0]][eaten[0][j - 1]]
        eaten[0][j] = eaten[0][j - 1] + 1
    else:
        break


for i in range(1, k):
    eaten[i][i] = 1
    dp[i][i] =  dp[i - 1][i - 1] + td[tp[i]][0] + 2 * i + 1
    for j in range(i + 1, k):
        if dp[i - 1][j - 1] == 0 or (eaten[i][j - 1] < len(td[tp[i]]) and 2 * i + 1 + td[tp[i]][0] < td[tp[i]][eaten[i][j - 1]]):
            if eaten[i][j - 1] < len(td[tp[i]]) and dp[i][j - 1] > 0:
                dp[i][j] = dp[i][j - 1] + td[tp[i]][eaten[i][j - 1]]
                eaten[i][j] = eaten[i][j - 1] + 1
        else:
            if dp[i - 1][j - 1] > 0:
                dp[i][j] = dp[i - 1][j - 1] + 2 * i + 1 + td[tp[i]][0]
                eaten[i][j] = 1

print(max([i[-1] for i in dp]))