Problem: There are N pieces of sushi. Each piece has two parameters: "kind of topping" t_i and "deliciousness" d_i. You are choosing K among these N pieces to eat. Your "satisfaction" here will be calculated as follows:    * The satisfaction is the sum of the "base total deliciousness" and the "variety bonus".   * The base total deliciousness is the sum of the deliciousness of the pieces you eat.   * The variety bonus is x*x, where x is the number of different kinds of toppings of the pieces you eat.  You want to have as much satisfaction as possible. Find this maximum satisfaction.

Submission attempts (in order):
#!/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 = [] 
    
    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)] 
    Y = [0] + [y for y in accumulate(B)] 

    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) 
            
            
            
    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()


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))


#!/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 = [] 
    
    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)] 
    Y = [0] + [y for y in accumulate(B)] 

    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) 
            
            
            
    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()


Verdicts: Runtime Error

Runtime Error

Accepted