Problem p03148 (python)

Fixed submission:
#!/usr/bin/env python3
import sys
from collections import defaultdict
from itertools import accumulate
import heapq

def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    sushi = defaultdict(list)
    for i in range(N):
        sushi[t[i]].append(d[i])

    A = [] # 各ネタでいちばんおいしい寿司のおいしさ
    B = [] # 各ネタで2番目以降においしい寿司のおいしさリスト
    
    for v in sushi.values():
        v.sort(reverse=True)
        A.append(v[0])
        B.extend(v[1:])
    A.sort(reverse=True)
    B.sort(reverse=True)

    X = [0] + [x for x in accumulate(A)] # X[i]: Aからうまい順にi個選んだ時のおいしさ合計
    Y = [0] + [y for y in accumulate(B)] # Y[j]: Bからうまい順にj個選んだ時のおいしさ合計

    res = 0
    for i in range(1, min(len(X), K + 1)):
        if K - i < len(Y):
            res = max(res, X[i] + Y[K - i] + i ** 2) 
            # B[:i]に含まれるネタが、A[i+1:]に含まれていた場合、実際のネタ数はiより多いが、
            # その場合はiを多くとったほうが、より美味しさが大きくなる。
            # そのためres[i]が答えになりえないため、無視して良い。
    print(res)
    return


# Generated by 1.1.4 https://github.com/kyuridenamida/atcoder-tools  (tips: You use the default template now. You can remove this line by using your custom template)
def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    K = int(next(tokens))  # type: int
    t = [int()] * (N)  # type: "List[int]" 
    d = [int()] * (N)  # type: "List[int]" 
    for i in range(N):
        t[i] = int(next(tokens))
        d[i] = int(next(tokens))
    solve(N, K, t, d)

if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
#!/usr/bin/env python3
import sys
from collections import defaultdict
from itertools import accumulate
import heapq

def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    sushi = defaultdict([])
    for i in range(N):
        sushi[t[i]].append(d[i])

    A = [] # 各ネタでいちばんおいしい寿司のおいしさ
    B = [] # 各ネタで2番目以降においしい寿司のおいしさリスト
    
    for v in sushi.values():
        v.sort(reverse=True)
        A.append(v[0])
        B.extend(v[1:])
    A.sort(reverse=True)
    B.sort(reverse=True)

    X = [0] + [x for x in accumulate(A)] # X[i]: Aからうまい順にi個選んだ時のおいしさ合計
    Y = [0] + [y for y in accumulate(B)] # Y[j]: Bからうまい順にj個選んだ時のおいしさ合計

    res = 0
    for i in range(1, min(len(X), K + 1)):
        if K - i < len(Y):
            res = max(res, X[i] + Y[K - i] + i ** 2) 
            # B[:i]に含まれるネタが、A[i+1:]に含まれていた場合、実際のネタ数はiより多いが、
            # その場合はiを多くとったほうが、より美味しさが大きくなる。
            # そのためres[i]が答えになりえないため、無視して良い。
    print(res)
    return


# Generated by 1.1.4 https://github.com/kyuridenamida/atcoder-tools  (tips: You use the default template now. You can remove this line by using your custom template)
def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    K = int(next(tokens))  # type: int
    t = [int()] * (N)  # type: "List[int]" 
    d = [int()] * (N)  # type: "List[int]" 
    for i in range(N):
        t[i] = int(next(tokens))
        d[i] = int(next(tokens))
    solve(N, K, t, d)

if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
#!/usr/bin/env python3
import sys
import heapq

def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    X = [(d[i], t[i]) for i in range(N)]
    X.sort(key=lambda x: x[1], reverse=True)

    Q = [] # うまい順にK個
    M = dict() # ネタごとに一番うまいやつ
    for i in range(N):
        if d[i] > M.get(t[i], 0):
            M[t[i]] = d[i]
        if len(Q) < K:
            heapq.heappush(Q, (d[i], t[i]))
        else:
            dq, tq = heapq.heappop(Q)
            if d[i] > dq:
                dq, tq = d[i], t[i]
            heapq.heappush(Q, (dq, tq))
    t_all = M.keys() # ネタ総数
    
    M = [(dm, tm) for tm, dm in M.items()]
    Q.sort(reverse=True)
    M.sort(reverse=True)

    T = [[0, set()]] # T[i=0~K]: うまい方からi個選び、のこりをMから選んだ時のおいしさと種類
    for i in range(K):
        t = T[-1][0] + Q[i][0]
        s = T[-1][1] | set([Q[i][1]])
        T.append([t, s])

    for i in range(K-1, -1, -1):
        if len(M) == 0:
            break
        while True:
            dm, tm = M.pop(0)
            if tm in T[i][1]:
                continue
            break
        else:
            break
        T[i][1] = T[i+1][1] | set([tm])
        T[i][0] = T[i+1][0] + dm - Q[i][0]

    res = 0
    for i in range(K+1):
        total = T[i][0] + len(T[i][1]) ** 2
        if total > res:
            res = total
    print(res)
    return


# Generated by 1.1.4 https://github.com/kyuridenamida/atcoder-tools  (tips: You use the default template now. You can remove this line by using your custom template)
def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    K = int(next(tokens))  # type: int
    t = [int()] * (N)  # type: "List[int]" 
    d = [int()] * (N)  # type: "List[int]" 
    for i in range(N):
        t[i] = int(next(tokens))
        d[i] = int(next(tokens))
    solve(N, K, t, d)

if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
sushi = [tuple(map(int, input().split())) for i in range(N)]
point = [0 for i in range(K+1)]

sushi.sort(lambda x: x[1])

eaten_sushi = [sushi[i] for i in range(K)]
eaten_variation = [s[0] for s in eaten_sushi]
variation_number = len(set(eaten_variation))
base = sum([s[1] for s in eaten_sushi])
bonus = variation_number ** 2
point[variation_number] = base + bonus

for i in range(K, N):
    if sushi[i][0] in eaten_variation:
        continue

    for j in range(K-1, 0, -1):
        if eaten_variation.count(eaten_sushi[j][0]) == 1:
            continue
        eaten_sushi.pop(j)
        eaten_sushi.append(sushi[i])
        eaten_variation = [s[0] for s in eaten_sushi]
        variation_number = len(set(eaten_variation))
        base = sum([s[1] for s in eaten_sushi])
        bonus = variation_number ** 2
        point[variation_number] = base + bonus
        break

print(max(point))