I found an interesting problem on https://atcoder.jp/contests/abc116/tasks/abc116_d:
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.

-----Constraints-----
 - 1 \leq K \leq N  \leq 10^5
 - 1 \leq t_i \leq N
 - 1 \leq d_i \leq 10^9
 - All values in input are integers.

-----Input-----
Input is given from Standard Input in the following format:
N K
t_1 d_1
t_2 d_2
.
.
.
t_N d_N

-----Output-----
Print the maximum satisfaction that you can obtain.

-----Sample Input-----
5 3
1 9
1 7
2 6
2 5
3 1

-----Sample Output-----
26

If you eat Sushi 1,2 and 3:
 - The base total deliciousness is 9+7+6=22.
 - The variety bonus is 2*2=4.
Thus, your satisfaction will be 26, which is optimal.

I tried it in Python, but could not do it. Can you solve it?
n,k=list(map(int,input().split()))
td=sorted([list(map(int,input().split())) for i in range(n)],reverse=True,key=lambda x:x[1])

ans=0
kl=dict()
for i in range(k):
    ans+=td[i][1]
    if td[i][0] in kl:
        kl[td[i][0]]+=1
    else:
        kl[td[i][0]]=1
l=len(kl)
ans+=(l**2)
ans_=ans
now=k-1
for i in range(k,n):
    if td[i][0] not in kl:
        #最小を求めるのをもっと早く
        while now>=0:
            if kl[td[now][0]]>1:
                mi=td[now]
                kl[mi[0]]-=1
                now-=1
                break
            now-=1
        if now==-1:
            break
        else:
            ans=ans+2*l+1-mi[1]+td[i][1]
            kl[td[i][0]]=1
            ans_=max(ans,ans_)
            l+=1

print(ans_)