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):
from math import pow

def compute_score(target_collection):
    keys = {neta_point[0] for neta_point in target_collection}
    points = [neta_point[1] for neta_point in target_collection]
    cardinality = len(keys)
    score = pow(cardinality, 2) + sum(points)
    return score

# exchange one record and compare
def optimize(full_collection, target_collection, target_collection_len, idx):
    while True:
        ex_record = full_collection[idx]
        # check on the idx-th record
        idx += 1
        # bump up for return value
        if ex_record in target_collection:
            continue # cardinality score will not increase
        else:
            for i in range(target_collection_len-1):
                if target_collection[i][0] == target_collection[i-1][0]:
                    target_collection[i] = ex_record
                    new_score = compute_score(target_collection)
                    return new_score, target_collection, idx
        
        return 0, target_collection, idx



def cal_score(N, K, records):
    sorted_recs = sorted(records, reverse=True, key=lambda x: x[1])
    target_collection = sorted_recs[0:K]
    score =  compute_score(target_collection)

    
    idx = K
    for i in range(K):
        new_score, new_target_collection, idx = optimize(sorted_recs, target_collection, K, idx)
        if new_score <= score: 
            break
        else:
            score = new_score
    return score
  
N, K = map(int, input().split())
records = [list(map(int, input().split())) for i in range(N)] 
score = cal_score(N, K, records)
print(score)


from math import pow



N, K = map(int, input().split())
records = [map(int, input().split()) for i in range(N)]


def compute_score(target_collection):
    keys = {neta_point[0] for neta_point in target_collection}
    points = [neta_point[1] for neta_point in target_collection]
    cardinality = len(keys)
    score = pow(cardinality, 2) + sum(points)
    return score

# exchange one record and compare
def optimize(full_collection, target_collection, target_collection_len, idx):
    while True:
        ex_record = full_collection[idx]
        # check on the idx-th record
        idx += 1
        # bump up for return value
        if ex_record in target_collection:
            continue # cardinality score will not increase
        else:
            for i in range(target_collection_len-1):
                if target_collection[i][0] == target_collection[i-1][0]:
                    target_collection[i] = ex_record
                    new_score = compute_score(target_collection)
                    return new_score, target_collection, idx
        
        return 0, target_collection, idx



def cal_score(N, K, records):
    sorted_recs = sorted(records, reverse=True, key=lambda x: x[1])
    target_collection = sorted_recs[0:K]
    score =  compute_score(target_collection)

    
    idx = K
    for i in range(K):
        new_score, new_target_collection, idx = optimize(sorted_recs, target_collection, K, idx)
        if new_score <= score: 
            break
        else:
            score = new_score
    return score
 
score =cal_score(N, K, record)
print(score)


from math import pow






def compute_score(target_collection):
    keys = {neta_point[0] for neta_point in target_collection}
    points = [neta_point[1] for neta_point in target_collection]
    cardinality = len(keys)
    score = pow(cardinality, 2) + sum(points)
    return score

# exchange one record and compare
def optimize(full_collection, target_collection, target_collection_len, idx):
    while True:
        ex_record = full_collection[idx]
        # check on the idx-th record
        idx += 1
        # bump up for return value
        if ex_record in target_collection:
            continue # cardinality score will not increase
        else:
            for i in range(target_collection_len-1):
                if target_collection[i][0] == target_collection[i-1][0]:
                    target_collection[i] = ex_record
                    new_score = compute_score(target_collection)
                    return new_score, target_collection, idx
        
        return 0, target_collection, idx



def cal_score(N, K, records):
    sorted_recs = sorted(records, reverse=True, key=lambda x: x[1])
    target_collection = sorted_recs[0:K]
    score =  compute_score(target_collection)

    
    idx = K
    for i in range(K):
        new_score, new_target_collection, idx = optimize(sorted_recs, target_collection, K, idx)
        if new_score <= score: 
            break
        else:
            score = new_score
    return score
  
N, K = map(int, input().split())
records = [tuple(map(int, input().split())) for i in range(N)] 
score = cal_score(N, K, records)
print(score)


from math import pow
from collections import deque, Counter
# from loguru import logger


def cal_score(N, K, records):
    sorted_recs = sorted(records, key=lambda x: -x[1])
    target_queue = deque(sorted_recs[0:K]) # only allow pop() action, so use dequeue
    
    compare_queue = deque(sorted_recs[K:])
    count = Counter()
    for neta_point in target_queue:
        count[neta_point[0]] += 1
    # logger.debug(count)

    cardinality = len(count.keys())
    sum_target_points = sum([n_p[1] for n_p in target_queue])
    score = pow(cardinality, 2) + sum_target_points
    score = int(score)
    new_scores = deque()
    new_scores.append(score)
    # start iteration and comparing gain and lost
    while compare_queue:
        comp = compare_queue.popleft()
        if comp[0] in count:
            continue
        while target_queue:
            old = target_queue.pop()
            if count[old[0]] < 2:
                # swap this neta will not increase cardinality but lower neta point
                continue
            else: # check if score will go up
                inc = 2 * cardinality + 1
                dec = old[1] - comp[1]
                delta = inc - dec
                # logger.debug(f"delta is {delta}")
                
                #     logger.debug("SOME")
                #     logger.debug(f"{old} out")
                #     logger.debug(f"{comp} in")
                #     cardinality += 1
                
                #     count[old[0]] -= 1
                #     count[comp[0]] +=1
                # else: #increase cardinality will bring more lost than gain
                #     # exit here
                #     return score
                new_score = new_scores[-1] + delta
                new_scores.append(new_score)

                cardinality += 1
                count[old[0]] -= 1
                count[comp[0]] +=1
    # both queues are drained
    return max(new_scores)
N, K = map(int, input().split())
records = [map(int, input().split()) for i in range(N)]
score = cal_score(N, K, records)
print(score)

from math import pow



N, K = map(int, input().split())
records = [map(int, input().split()) for i in range(N)]


def compute_score(target_collection):
    keys = {neta_point[0] for neta_point in target_collection}
    points = [neta_point[1] for neta_point in target_collection]
    cardinality = len(keys)
    score = pow(cardinality, 2) + sum(points)
    return score

# exchange one record and compare
def optimize(full_collection, target_collection, target_collection_len, idx):
    while True:
        ex_record = full_collection[idx]
        # check on the idx-th record
        idx += 1
        # bump up for return value
        if ex_record in target_collection:
            continue # cardinality score will not increase
        else:
            for i in range(target_collection_len-1):
                if target_collection[i][0] == target_collection[i-1][0]:
                    target_collection[i] = ex_record
                    new_score = compute_score(target_collection)
                    return new_score, target_collection, idx
        
        return 0, target_collection, idx



def cal_score(N, K, records):
    sorted_recs = sorted(records, reverse=True, key=lambda x: x[1])
    target_collection = sorted_recs[0:K]
    score =  compute_score(target_collection)

    
    idx = K
    for i in range(K):
        new_score, new_target_collection, idx = optimize(sorted_recs, target_collection, K, idx)
        if new_score <= score: 
            break
        else:
            score = new_score
    return score
 
score = cal_score(N, K, records)
print(score)


from math import pow
from collections import deque, Counter
import sys
stdin = sys.stdin
def li(): return map(int, stdin.readline().split())

def cal_score(N, K, records):
    sorted_recs = sorted(records, key=lambda x: -x[1])
    target_queue = deque(sorted_recs[0:K]) # only allow pop() action, so use dequeue
    compare_queue = deque(sorted_recs[K:])
    distinct_set = set()
    dup_queue = deque()
    score = 0
    for neta_point in target_queue:
        score += neta_point[1]
        if neta_point[0] not in distinct_set:
            distinct_set.add(neta_point[0])
        else:
            dup_queue.append(neta_point)

    cardinality = len(distinct_set)
    score += pow(cardinality, 2)
    score = int(score)
    maybe_score = score
    # start iteration and comparing gain and lost
    while compare_queue:
        comp = compare_queue.popleft()
        if comp[0] in distinct_set:
            continue
        elif len(dup_queue) < 1:
            break
        else:
            old = dup_queue.pop()
            inc = 2 * len(distinct_set) + 1
            dec = old[1] - comp[1]
            delta = inc - dec
            maybe_score = maybe_score + delta
            score = max(score, maybe_score)

            distinct_set.add(comp[0])
    # both queues are drained
    return score


N, K = li()
records = [tuple(li()) for i in range(N)]
score = cal_score(N, K, records)
print(score)


Verdicts: Runtime Error

Runtime Error

Runtime Error

Runtime Error

Runtime Error

Accepted