Problem p03148 (python)

Fixed submission:
def solve(string):
    n, k, *td = map(int, string.split())
    td = sorted([(t, d) for t, d in zip(td[::2], td[1::2])], key=lambda x: x[1], reverse=True)
    max_variety = len(set([c[0] for c in td]))
    choice = td[:k]
    count = [0] * n
    r = []
    for _c in choice:
        if count[_c[0] - 1] > 0:
            r.append(_c[1])
        count[_c[0] - 1] += 1
    variety = set([c[0] for c in choice])
    ans = [len(variety)**2 + sum([c[1] for c in choice])]
    i_append = k
    while len(variety) < min(max_variety, k):
        while td[i_append][0] in variety:
            i_append += 1
        variety.add(td[i_append][0])
        ans.append(ans[-1] - r.pop() + td[i_append][1] + 2 * len(variety) - 1)
    return str(max(ans))


if __name__ == '__main__':
    n, m = map(int, input().split())
    print(solve('{} {}\n'.format(n, m) + '\n'.join([input() for _ in range(n)])))

Buggy submission (verdict: Runtime Error):
def solve(string):
    n, k, *td = map(int, string.split())
    td = sorted([(t, d) for t, d in zip(td[::2], td[1::2])], key=lambda x: x[1], reverse=True)
    max_variety = len(set([c[0] for c in td]))
    choice = td[:k]
    r = []
    variety = set([])
    ans = 0
    for _c in choice:
        ans += _c[1]
        if _c[0] in variety:
            r.append(_c[1])
        else:
            variety.add(_c[0])
    ans += len(variety)**2
    i = k
    while len(variety) < min(max_variety, k):
        while td[i][0] in variety:
            i += 1
        variety.add(td[i][0])
        ans.append(ans[-1] - r.pop() + td[i][1] + 2 * len(variety) - 1)
    return str(max(ans))


if __name__ == '__main__':
    n, m = map(int, input().split())
    print(solve('{} {}\n'.format(n, m) + '\n'.join([input() for _ in range(n)])))

Buggy submission (verdict: Runtime Error):
from operator import itemgetter
from collections import Counter


def solve(string):
    n, k, *td = map(int, string.split())
    dic = {i: [] for i in range(1, n + 1)}
    td = sorted([(t, d) for t, d in zip(td[::2], td[1::2])], key=itemgetter(1), reverse=True)
    for t, d in td:
        dic[t].append(d)
    base = td[:k]
    counts = Counter([_b[0] for _b in base])
    v = set(counts.keys())
    start_v = len(v)
    max_v = min(len(set(td[::2])), k)

    candidates = []
    for _td in td[k:]:
        if _td[0] not in v:
            candidates.append(_td)
            v.add(_td[0])

    ans = sum([_b[1] for _b in base]) + start_v**2
    index = -1
    tmp = ans
    for i in range(max_v - start_v):
        while counts[base[index][0]] == 1:
            index -= 1
        _t, _d = base[index]
        counts[_t] -= 1
        tmp = tmp - _d + candidates[i][1] + 2 * (start_v + i + 1) - 1
        ans = max(tmp, ans)

    return str(ans)


if __name__ == '__main__':
    n, m = map(int, input().split())
    print(solve('{} {}\n'.format(n, m) + '\n'.join([input() for _ in range(n)])))

Buggy submission (verdict: Runtime Error):
from copy import deepcopy


def solve(string):
    n, k, *td = map(int, string.split())
    dic = {i: [] for i in range(1, n + 1)}
    for t, d in zip(td[::2], td[1::2]):
        dic[t].append(d)
    dic = {_k: sorted(dic[_k], reverse=True) for _k in dic.keys()}

    def dfs(variety, b, use):
        if sum(use) + n == k:
            return len(variety)**2 + b
        tmp = []
        for _u, _k in zip(use, dic.keys()):
            if _u < len(dic[_k]) - 1:
                u = deepcopy(use)
                v = deepcopy(variety)
                u[_k - 1] += 1
                v.add(_k)
                tmp.append(dfs(v, b + dic[_k][u[_k - 1]], u))
        return max(tmp)

    return str(dfs(set([]), 0, [-1] * n))


if __name__ == '__main__':
    n, m = map(int, input().split())
    print(solve('{} {}\n'.format(n, m) + '\n'.join([input() for _ in range(n)])))