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


INPUT1 = [int(n) for n in input().split()]

N, K = INPUT1[0], INPUT1[1]
Parameters = []
Score, Test = 0, 0
flag1, flag2 = 0, 0


for i in range(N):
    INPUT2 = [int(m) for m in input().split()]
    Parameters.append((INPUT2[0], INPUT2[1]))

Parameters = sorted(Parameters, key=lambda x: x[1], reverse=True)

def SumScore(List, K):
    taste = 0
    kind = []
    for i in range(K):
        taste += List[i][1]
        kind.append(List[i][0])
    return taste + len(set(kind))**2

def SearchNew(List, K):
    kind = []
    newkind = -1
    for i in range(K):
        kind.append(List[i][0])
    kindset = set(kind)
    for i in range(len(List)-K):
        if List[K+i][0] not in kindset:
            newkind = K+i
            break
    return newkind

def SearchDouble(List, K):
    result = -1
    for i in range(K):
        leftkind = []
        for j in range(K):
            leftkind.append(List[j][0])
        testkind = List[K-1-i][0]
        leftkind.pop(K-1-i)
        if testkind in leftkind:
            result = K-1-i
            break
    return result

Score = SumScore(Parameters, K)

for t in Parameters:
    i = SearchNew(Parameters, K)
    j = SearchDouble(Parameters, K)
    Parameters[i], Parameters[j] = Parameters[j], Parameters[i]
    Test = SumScore(Parameters, K)
    if Score <= Test:
        Score = Test
    print(i, j, t)

print(Score)




INPUT1 = [int(n) for n in input().split()]

N, K = INPUT1[0], INPUT1[1]
Parameters = []
Score, Test = 0, 0
flag1, flag2 = 0, 0
kind = []
kindnumber = 0

ParametersAppend = Parameters.append
for i in range(N):
    INPUT2 = [int(m) for m in input().split()]
    ParametersAppend([INPUT2[0], INPUT2[1]])

Parameters = sorted(Parameters, key=lambda x: x[1], reverse=True)

def binanrySearch(kind, target):
    testkind = kind[:]
    testkind.sort()
    left, right = 0, len(testkind)-1
    while(right >= left):
        mid = int((right + left) / 2)
        if testkind[mid] == target:
            return True
        elif testkind[mid] > target:
            right = mid - 1
        elif testkind[mid] < target:
            left = mid + 1
    return False

def SumScore(List, K, kindnumber):
    taste = 0
    for i in range(K):
        taste += List[i][1]
    return taste + kindnumber * kindnumber

def SearchNew(List, K, kind):
    newkind = -1
    for i in range(len(List)-K):
        if binanrySearch(kind, List[K+i][0]) == False:
            newkind = K+i
            break
    return newkind

def SearchDouble(List, K, kind):
    result = -1
    testkind = kind[:]
    testkindPop = testkind.pop
    for i in range(K):
        testkindPop(K-1-i)
        if binanrySearch(testkind, List[K-1-i][0]) == True:
            result = K-1-i
            break
    return result

kindAppend = kind.append

for i in range(K):
    kindAppend(Parameters[i][0])
kindnumber = len(set(kind))
Score = SumScore(Parameters, K, kindnumber)
kind.clear()

for t in Parameters:
    for i in range(K):
        kindAppend(Parameters[i][0])
    kindnumber = len(set(kind))
    i = SearchNew(Parameters, K, kind)
    j = SearchDouble(Parameters, K, kind)
    Parameters[i], Parameters[j] = Parameters[j], Parameters[i]
    if i != -1 and j != -1:
        kindnumber += 1
    else:
        break
    Test = Score - Parameters[i][1] + Parameters[j][1] + kindnumber**2
    if Score <= Test:
        Score = Test
    kind.clear()

print(Score)




from sys import stdin
#from numpy import sum

input = stdin.readline

N, K = list(map(int, input().split()))
List1= []
kind = set()
point1, point2 = 0, 0

INPUT = [tuple(map(int, input().split())) for _ in range(N)]

INPUT.sort(key = lambda x: -x[1])

for ti, di in INPUT[:K]:
    point1 += di
    if ti in kind:
        List1.append(di)
    kind.add(ti)

point2 = len(kind)

Score = point1 + point2**2
Test = Score

List1POP = List1.pop

for ti, di in INPUT[K:]:
    if 0 >= len(List1):
        break
    if ti in kind:
        continue
    Test = Test - List1POP() + di + 2 * point2 + 1
    point2 += 1
    kind.add(ti)
    Score = max(Score, Test)

print(Score)


Verdicts: Wrong Answer

Wrong Answer

Accepted