Buggy python submission (verdict: Time Limit Exceeded):
# -*- coding: utf-8 -*-
import sys
import math
from bisect import bisect_left
from bisect import bisect_right
from collections import defaultdict
from heapq import heappop, heappush
import itertools
from decimal import *

input = sys.stdin.readline

def inputInt(): return int(input())
def inputMap(): return map(int, input().split())
def inputList(): return list(map(int, input().split()))
def inputStr(): return input()[:-1]

inf = float('inf')
mod = 1000000007

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-

def main():
	N,K = inputMap()
	sushi_mi = []
	sushi_ki = []
	syu = [[] for i in range(N)]
	for i in range(N):
		t,d = inputMap()
		t -= 1
		syu[t].append(d)

	for i in syu:
		if len(i) > 0:
			i.sort()
			i = i[::-1]
			for j,val in enumerate(i):
				if j == 0:
					sushi_mi.append(val)
				else:
					sushi_ki.append(val)

	sushi_mi.sort()
	sushi_mi = sushi_mi[::-1]
	sushi_ki.sort()
	sushi_ki = sushi_ki[::-1]

	#print(sushi_mi)
	#print(sushi_ki)

	ans = 0
	for i in range(1, min(len(sushi_mi), K)+1):
		#print(min(len(sushi_mi), K+1)+1)
		x = 0
		tmp = 0
		if K - i > len(sushi_ki):
			continue

		for j in range(0, i):
			tmp += sushi_mi[j]
			x += 1

		for j in range(0, K - i):
			tmp += sushi_ki[j]

		oishi = tmp + x * x
		#print(oishi)
		if ans < oishi:
			ans = oishi
	print(ans)

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
if __name__ == "__main__":
	main()


Fixed submission:
# -*- coding: utf-8 -*-
import sys
import math
from bisect import bisect_left
from bisect import bisect_right
from collections import defaultdict
from heapq import heappop, heappush
import itertools
from decimal import *

input = sys.stdin.readline

def inputInt(): return int(input())
def inputMap(): return map(int, input().split())
def inputList(): return list(map(int, input().split()))
def inputStr(): return input()[:-1]

inf = float('inf')
mod = 1000000007

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-

def main():
	N,K = inputMap()
	sushi_mii = []
	sushi_kii = []
	syu = [[] for i in range(N)]
	for i in range(N):
		t,d = inputMap()
		t -= 1
		syu[t].append(d)

	for i in syu:
		if len(i) > 0:
			i.sort()
			i = i[::-1]
			for j,val in enumerate(i):
				if j == 0:
					sushi_mii.append(val)
				else:
					sushi_kii.append(val)

	sushi_mii.sort()
	sushi_mii = sushi_mii[::-1]
	sushi_kii.sort()
	sushi_kii = sushi_kii[::-1]

	sushi_mi = []
	sushi_ki = []
	for i,val in enumerate(sushi_mii):
		if i == 0:
			sushi_mi.append(val)
		else:
			tmp = sushi_mi[i-1]
			sushi_mi.append(tmp + val)

	for i,val in enumerate(sushi_kii):
		if i == 0:
			sushi_ki.append(val)
		else:
			tmp = sushi_ki[i-1]
			sushi_ki.append(tmp + val)

	ans = 0
	for i in range(1, min(len(sushi_mi), K)+1):
		x = 0
		tmp = 0
		if K - i > len(sushi_ki):
			continue

		tmp += sushi_mi[i-1]
		x = i

		if K - i > 0:
			tmp += sushi_ki[K - i-1]

		oishi = tmp + x * x
		if ans < oishi:
			ans = oishi
	print(ans)

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
if __name__ == "__main__":
	main()

Buggy python submission (verdict: Runtime Error):
# -*- coding: utf-8 -*-
import sys
import math
from bisect import bisect_left
from bisect import bisect_right
from collections import defaultdict
from heapq import heappop, heappush
import itertools
from decimal import *

input = sys.stdin.readline

def inputInt(): return int(input())
def inputMap(): return map(int, input().split())
def inputList(): return list(map(int, input().split()))
def inputStr(): return input()[:-1]

inf = float('inf')
mod = 1000000007

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-

def main():
	N,K = inputMap()
	sushi_mi = []
	sushi_ki = []
	syu = {}
	for i in range(N):
		t,d = inputMap()
		if t in syu:
			sushi_ki.append(d)
		else:
			syu[t] = 1
			sushi_mi.append(d)

	sushi_mi.sort()
	sushi_mi = sushi_mi[::-1]
	sushi_ki.sort()
	sushi_ki = sushi_ki[::-1]

	ans = 0
	for i in range(1,K+1):
		tmp = 0
		for j in range(0,i):
			tmp += sushi_mi[j]
		for j in range(0,K-i):
			tmp += sushi_ki[j]

		tmp += i*i
		if tmp > ans:
			ans = tmp

	print(ans)

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
if __name__ == "__main__":
	main()


Fixed submission:
# -*- coding: utf-8 -*-
import sys
import math
from bisect import bisect_left
from bisect import bisect_right
from collections import defaultdict
from heapq import heappop, heappush
import itertools
from decimal import *

input = sys.stdin.readline

def inputInt(): return int(input())
def inputMap(): return map(int, input().split())
def inputList(): return list(map(int, input().split()))
def inputStr(): return input()[:-1]

inf = float('inf')
mod = 1000000007

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-

def main():
	N,K = inputMap()
	sushi_mii = []
	sushi_kii = []
	syu = [[] for i in range(N)]
	for i in range(N):
		t,d = inputMap()
		t -= 1
		syu[t].append(d)

	for i in syu:
		if len(i) > 0:
			i.sort()
			i = i[::-1]
			for j,val in enumerate(i):
				if j == 0:
					sushi_mii.append(val)
				else:
					sushi_kii.append(val)

	sushi_mii.sort()
	sushi_mii = sushi_mii[::-1]
	sushi_kii.sort()
	sushi_kii = sushi_kii[::-1]

	sushi_mi = []
	sushi_ki = []
	for i,val in enumerate(sushi_mii):
		if i == 0:
			sushi_mi.append(val)
		else:
			tmp = sushi_mi[i-1]
			sushi_mi.append(tmp + val)

	for i,val in enumerate(sushi_kii):
		if i == 0:
			sushi_ki.append(val)
		else:
			tmp = sushi_ki[i-1]
			sushi_ki.append(tmp + val)

	ans = 0
	for i in range(1, min(len(sushi_mi), K)+1):
		x = 0
		tmp = 0
		if K - i > len(sushi_ki):
			continue

		tmp += sushi_mi[i-1]
		x = i

		if K - i > 0:
			tmp += sushi_ki[K - i-1]

		oishi = tmp + x * x
		if ans < oishi:
			ans = oishi
	print(ans)

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
if __name__ == "__main__":
	main()

Buggy python submission (verdict: Time Limit Exceeded):
# -*- coding: utf-8 -*-
import sys
from bisect import bisect_right

def main():
	N, K = map(int, input().split(" "))
	osushi = []
	for _ in  range(N):
		t, d = map(int, input().split(" "))
		osushi.append((d, t))
		
	osushi.sort()
	osushi.reverse()
	
	osushi_list = []
	syu_total = 0
	syu_list = []
	man = 0
	syu = 0
	for i, val in enumerate(osushi):
		d, t = val;
		if not t in syu_list:
			syu_total += 1
			syu_list.append(t)
			t = 1
		else:
			t = 0
		osushi_list.append((d, t))
		if i < K:
			man += d
			if t == 1:
				syu += 1
	
	ans = man + (syu ** 2)
	tmp1 = sorted(osushi_list, key = lambda k:k[1], reverse=True)
	tmp2 = sorted(osushi_list, key = lambda k:k[1])
	for syu_mokuhyou in range(syu+1, syu_total+1):
		man = 0
		syu_local = 0
		count = 0
		for i, val in enumerate(tmp1):
			d, t = val
			syu_local += 1
			man += d
			count += 1
			if syu_local == syu_mokuhyou or count >= K:
				break
					
		for i, val in enumerate(tmp2):
			if count >= K:
				break
			d, t = val
			man += d
			count += 1
		tmp = man + (syu_local ** 2)
		if tmp > ans:
			ans = tmp
		
	print(ans)
	
	
if __name__ == "__main__":
	main()


Fixed submission:
# -*- coding: utf-8 -*-
import sys
import math
from bisect import bisect_left
from bisect import bisect_right
from collections import defaultdict
from heapq import heappop, heappush
import itertools
from decimal import *

input = sys.stdin.readline

def inputInt(): return int(input())
def inputMap(): return map(int, input().split())
def inputList(): return list(map(int, input().split()))
def inputStr(): return input()[:-1]

inf = float('inf')
mod = 1000000007

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-

def main():
	N,K = inputMap()
	sushi_mii = []
	sushi_kii = []
	syu = [[] for i in range(N)]
	for i in range(N):
		t,d = inputMap()
		t -= 1
		syu[t].append(d)

	for i in syu:
		if len(i) > 0:
			i.sort()
			i = i[::-1]
			for j,val in enumerate(i):
				if j == 0:
					sushi_mii.append(val)
				else:
					sushi_kii.append(val)

	sushi_mii.sort()
	sushi_mii = sushi_mii[::-1]
	sushi_kii.sort()
	sushi_kii = sushi_kii[::-1]

	sushi_mi = []
	sushi_ki = []
	for i,val in enumerate(sushi_mii):
		if i == 0:
			sushi_mi.append(val)
		else:
			tmp = sushi_mi[i-1]
			sushi_mi.append(tmp + val)

	for i,val in enumerate(sushi_kii):
		if i == 0:
			sushi_ki.append(val)
		else:
			tmp = sushi_ki[i-1]
			sushi_ki.append(tmp + val)

	ans = 0
	for i in range(1, min(len(sushi_mi), K)+1):
		x = 0
		tmp = 0
		if K - i > len(sushi_ki):
			continue

		tmp += sushi_mi[i-1]
		x = i

		if K - i > 0:
			tmp += sushi_ki[K - i-1]

		oishi = tmp + x * x
		if ans < oishi:
			ans = oishi
	print(ans)

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
if __name__ == "__main__":
	main()

Buggy python submission (verdict: Time Limit Exceeded):
# -*- coding: utf-8 -*-
import sys
from bisect import bisect_right

def main():
	N, K = map(int, input().split(" "))
	osushi = []
	for _ in  range(N):
		t, d = map(int, input().split(" "))
		osushi.append((d, t))
		
	osushi.sort()
	osushi.reverse()
	
	osushi_list = []
	syu_total = 0
	syu_list = []
	man = 0
	syu = 0
	for i, val in enumerate(osushi):
		d, t = val;
		if not t in syu_list:
			syu_total += 1
			syu_list.append(t)
			t = 1
		else:
			t = 0
		osushi_list.append((d, t))
		if i < K:
			man += d
			if t == 1:
				syu += 1
	
	ans = man + (syu ** 2)
	for syu_mokuhyou in range(syu+1, syu_total+1):
		man = 0
		syu_local = 0
		count = 0
		tmp = sorted(osushi_list, key = lambda k:k[1], reverse=True)
		for i, val in enumerate(tmp):
			d, t = val
			syu_local += 1
			man += d
			count += 1
			if syu_local == syu_mokuhyou or count >= K:
				break
					
		tmp = sorted(osushi_list, key = lambda k:k[1])
		for i, val in enumerate(tmp):
			if count >= K:
				break
			d, t = val
			man += d
			count += 1
		tmp = man + (syu_local ** 2)
		if tmp > ans:
			ans = tmp
		
	print(ans)
	
	
if __name__ == "__main__":
	main()


Fixed submission:
# -*- coding: utf-8 -*-
import sys
import math
from bisect import bisect_left
from bisect import bisect_right
from collections import defaultdict
from heapq import heappop, heappush
import itertools
from decimal import *

input = sys.stdin.readline

def inputInt(): return int(input())
def inputMap(): return map(int, input().split())
def inputList(): return list(map(int, input().split()))
def inputStr(): return input()[:-1]

inf = float('inf')
mod = 1000000007

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-

def main():
	N,K = inputMap()
	sushi_mii = []
	sushi_kii = []
	syu = [[] for i in range(N)]
	for i in range(N):
		t,d = inputMap()
		t -= 1
		syu[t].append(d)

	for i in syu:
		if len(i) > 0:
			i.sort()
			i = i[::-1]
			for j,val in enumerate(i):
				if j == 0:
					sushi_mii.append(val)
				else:
					sushi_kii.append(val)

	sushi_mii.sort()
	sushi_mii = sushi_mii[::-1]
	sushi_kii.sort()
	sushi_kii = sushi_kii[::-1]

	sushi_mi = []
	sushi_ki = []
	for i,val in enumerate(sushi_mii):
		if i == 0:
			sushi_mi.append(val)
		else:
			tmp = sushi_mi[i-1]
			sushi_mi.append(tmp + val)

	for i,val in enumerate(sushi_kii):
		if i == 0:
			sushi_ki.append(val)
		else:
			tmp = sushi_ki[i-1]
			sushi_ki.append(tmp + val)

	ans = 0
	for i in range(1, min(len(sushi_mi), K)+1):
		x = 0
		tmp = 0
		if K - i > len(sushi_ki):
			continue

		tmp += sushi_mi[i-1]
		x = i

		if K - i > 0:
			tmp += sushi_ki[K - i-1]

		oishi = tmp + x * x
		if ans < oishi:
			ans = oishi
	print(ans)

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
if __name__ == "__main__":
	main()

Buggy python submission (verdict: Time Limit Exceeded):
# -*- coding: utf-8 -*-
import sys
from bisect import bisect_right

def main():
	N, K = map(int, input().split(" "))
	osushi = []
	for _ in  range(N):
		t, d = map(int, input().split(" "))
		osushi.append((d, t))
		
	osushi.sort()
	osushi.reverse()
	
	osushi_list = []
	syu_total = 0
	syu_list = []
	man = 0
	syu = 0
	for i, val in enumerate(osushi):
		d, t = val;
		if not t in syu_list:
			syu_total += 1
			syu_list.append(t)
			t = 1
		else:
			t = 0
		osushi_list.append((d, t))
		if i < K:
			man += d
			if t == 1:
				syu += 1
	
	ans = man + (syu ** 2)
	for syu_mokuhyou in range(syu+1, syu_total+1):
		man = 0
		syu_local = 0
		count = 0
		tmp = sorted(osushi_list, key = lambda k:k[1], reverse=True)
		for i, val in enumerate(tmp):
			d, t = val
			syu_local += 1
			man += d
			count += 1
			if syu_local == syu_mokuhyou or count >= K:
				break
					
		tmp = sorted(osushi_list, key = lambda k:k[1])
		for i, val in enumerate(tmp):
			if count >= K:
				break
			d, t = val
			man += d
			count += 1
		tmp = man + (syu_local ** 2)
		if tmp > ans:
			ans = tmp
		
	print(ans)
	
	
if __name__ == "__main__":
	main()


Fixed submission:
# -*- coding: utf-8 -*-
import sys
import math
from bisect import bisect_left
from bisect import bisect_right
from collections import defaultdict
from heapq import heappop, heappush
import itertools
from decimal import *

input = sys.stdin.readline

def inputInt(): return int(input())
def inputMap(): return map(int, input().split())
def inputList(): return list(map(int, input().split()))
def inputStr(): return input()[:-1]

inf = float('inf')
mod = 1000000007

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-

def main():
	N,K = inputMap()
	sushi_mii = []
	sushi_kii = []
	syu = [[] for i in range(N)]
	for i in range(N):
		t,d = inputMap()
		t -= 1
		syu[t].append(d)

	for i in syu:
		if len(i) > 0:
			i.sort()
			i = i[::-1]
			for j,val in enumerate(i):
				if j == 0:
					sushi_mii.append(val)
				else:
					sushi_kii.append(val)

	sushi_mii.sort()
	sushi_mii = sushi_mii[::-1]
	sushi_kii.sort()
	sushi_kii = sushi_kii[::-1]

	sushi_mi = []
	sushi_ki = []
	for i,val in enumerate(sushi_mii):
		if i == 0:
			sushi_mi.append(val)
		else:
			tmp = sushi_mi[i-1]
			sushi_mi.append(tmp + val)

	for i,val in enumerate(sushi_kii):
		if i == 0:
			sushi_ki.append(val)
		else:
			tmp = sushi_ki[i-1]
			sushi_ki.append(tmp + val)

	ans = 0
	for i in range(1, min(len(sushi_mi), K)+1):
		x = 0
		tmp = 0
		if K - i > len(sushi_ki):
			continue

		tmp += sushi_mi[i-1]
		x = i

		if K - i > 0:
			tmp += sushi_ki[K - i-1]

		oishi = tmp + x * x
		if ans < oishi:
			ans = oishi
	print(ans)

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
if __name__ == "__main__":
	main()

Buggy python submission (verdict: Time Limit Exceeded):
# -*- coding: utf-8 -*-
import sys
from bisect import bisect_right

def main():
	N, K = map(int, input().split(" "))
	osushi = []
	for _ in  range(N):
		t, d = map(int, input().split(" "))
		osushi.append((d, t))
		
	osushi.sort()
	osushi.reverse()
	
	osushi_list = []
	syu_total = 0
	syu_list = []
	man = 0
	syu = 0
	for i, val in enumerate(osushi):
		d, t = val;
		if not t in syu_list:
			syu_total += 1
			syu_list.append(t)
			t = 1
		else:
			t = 0
		osushi_list.append((d, t))
		if i < K:
			man += d
			if t == 1:
				syu += 1
	
	ans = man + (syu ** 2)
	for syu_mokuhyou in range(syu+1, syu_total+1):
		man = 0
		syu_local = 0
		count = 0
		for i, val in enumerate(osushi_list):
			d, t = val
			if t == 1:
				syu_local += 1
				man += d
				count += 1
				if syu_local == syu_mokuhyou or count >= K:
					break
					
		for i, val in enumerate(osushi_list):
			if count >= K:
				break
			d, t = val
			if t == 0:
				man += d
				count += 1
		tmp = man + (syu_local ** 2)
		if tmp > ans:
			ans = tmp
		
	print(ans)
	
	
if __name__ == "__main__":
	main()


Fixed submission:
# -*- coding: utf-8 -*-
import sys
import math
from bisect import bisect_left
from bisect import bisect_right
from collections import defaultdict
from heapq import heappop, heappush
import itertools
from decimal import *

input = sys.stdin.readline

def inputInt(): return int(input())
def inputMap(): return map(int, input().split())
def inputList(): return list(map(int, input().split()))
def inputStr(): return input()[:-1]

inf = float('inf')
mod = 1000000007

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-

def main():
	N,K = inputMap()
	sushi_mii = []
	sushi_kii = []
	syu = [[] for i in range(N)]
	for i in range(N):
		t,d = inputMap()
		t -= 1
		syu[t].append(d)

	for i in syu:
		if len(i) > 0:
			i.sort()
			i = i[::-1]
			for j,val in enumerate(i):
				if j == 0:
					sushi_mii.append(val)
				else:
					sushi_kii.append(val)

	sushi_mii.sort()
	sushi_mii = sushi_mii[::-1]
	sushi_kii.sort()
	sushi_kii = sushi_kii[::-1]

	sushi_mi = []
	sushi_ki = []
	for i,val in enumerate(sushi_mii):
		if i == 0:
			sushi_mi.append(val)
		else:
			tmp = sushi_mi[i-1]
			sushi_mi.append(tmp + val)

	for i,val in enumerate(sushi_kii):
		if i == 0:
			sushi_ki.append(val)
		else:
			tmp = sushi_ki[i-1]
			sushi_ki.append(tmp + val)

	ans = 0
	for i in range(1, min(len(sushi_mi), K)+1):
		x = 0
		tmp = 0
		if K - i > len(sushi_ki):
			continue

		tmp += sushi_mi[i-1]
		x = i

		if K - i > 0:
			tmp += sushi_ki[K - i-1]

		oishi = tmp + x * x
		if ans < oishi:
			ans = oishi
	print(ans)

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
if __name__ == "__main__":
	main()

Buggy python submission (verdict: Time Limit Exceeded):
# -*- coding: utf-8 -*-
import sys
from bisect import bisect_right

def main():
	N, K = map(int, input().split(" "))
	osushi = []
	for _ in  range(N):
		t, d = map(int, input().split(" "))
		osushi.append((d, t))
		
	osushi.sort()
	osushi.reverse()
	
	osushi_list = []
	syu_total = 0
	syu_list = []
	for i, val in enumerate(osushi):
		d, t = val;
		if not t in syu_list:
			syu_total += 1
			syu_list.append(t)
			t = 1
		else:
			t = 0
		osushi_list.append((d, t))
	
	man = 0
	syu = 0
	for i, val in enumerate(osushi_list):
		if i >= K:
			break
		d, t = val
		if t == 1:
			syu += 1
		man += d
	
	ans = man + (syu ** 2)
	for syu_mokuhyou in range(syu+1, syu_total+1):
		man = 0
		syu_local = 0
		count = 0
		for i, val in enumerate(osushi_list):
			d, t = val
			if t == 1:
				syu_local += 1
				man += d
				count += 1
				if syu_local == syu_mokuhyou or count >= K:
					break
					
		for i, val in enumerate(osushi_list):
			if count >= K:
				break
			d, t = val
			if t == 0:
				man += d
				count += 1
		tmp = man + (syu_local ** 2)
		if tmp > ans:
			ans = tmp
		
	print(ans)
	
	
if __name__ == "__main__":
	main()


Fixed submission:
# -*- coding: utf-8 -*-
import sys
import math
from bisect import bisect_left
from bisect import bisect_right
from collections import defaultdict
from heapq import heappop, heappush
import itertools
from decimal import *

input = sys.stdin.readline

def inputInt(): return int(input())
def inputMap(): return map(int, input().split())
def inputList(): return list(map(int, input().split()))
def inputStr(): return input()[:-1]

inf = float('inf')
mod = 1000000007

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-

def main():
	N,K = inputMap()
	sushi_mii = []
	sushi_kii = []
	syu = [[] for i in range(N)]
	for i in range(N):
		t,d = inputMap()
		t -= 1
		syu[t].append(d)

	for i in syu:
		if len(i) > 0:
			i.sort()
			i = i[::-1]
			for j,val in enumerate(i):
				if j == 0:
					sushi_mii.append(val)
				else:
					sushi_kii.append(val)

	sushi_mii.sort()
	sushi_mii = sushi_mii[::-1]
	sushi_kii.sort()
	sushi_kii = sushi_kii[::-1]

	sushi_mi = []
	sushi_ki = []
	for i,val in enumerate(sushi_mii):
		if i == 0:
			sushi_mi.append(val)
		else:
			tmp = sushi_mi[i-1]
			sushi_mi.append(tmp + val)

	for i,val in enumerate(sushi_kii):
		if i == 0:
			sushi_ki.append(val)
		else:
			tmp = sushi_ki[i-1]
			sushi_ki.append(tmp + val)

	ans = 0
	for i in range(1, min(len(sushi_mi), K)+1):
		x = 0
		tmp = 0
		if K - i > len(sushi_ki):
			continue

		tmp += sushi_mi[i-1]
		x = i

		if K - i > 0:
			tmp += sushi_ki[K - i-1]

		oishi = tmp + x * x
		if ans < oishi:
			ans = oishi
	print(ans)

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
if __name__ == "__main__":
	main()

Buggy python submission (verdict: Time Limit Exceeded):
# -*- coding: utf-8 -*-
import sys
from bisect import bisect_right

def main():
	N, K = map(int, input().split(" "))
	osushi = []
	for _ in  range(N):
		t, d = map(int, input().split(" "))
		osushi.append((d, t))
		
	osushi.sort()
	osushi.reverse()
	
	osushi_list = []
	syu_total = 0
	syu_list = []
	for i, val in enumerate(osushi):
		d, t = val;
		if not t in syu_list:
			syu_total += 1
			syu_list.append(t)
			t = 1
		else:
			t = 0
		osushi_list.append((d, t))
	
	man = 0
	syu = 0
	for i, val in enumerate(osushi_list):
		if i >= K:
			break
		d, t = val
		if t == 1:
			syu += 1
		man += d
	
	ans = man + (syu ** 2)
	for syu_mokuhyou in range(syu+1, syu_total+1):
		man = 0
		syu_local = 0
		count = 0
		for i, val in enumerate(osushi_list):
			d, t = val
			if t == 1:
				syu_local += 1
				man += d
				count += 1
				if syu_local == syu_mokuhyou or count >= K:
					break
					
		for i, val in enumerate(osushi_list):
			if count >= K:
				break
			d, t = val
			if t == 0:
				man += d
				count += 1
		tmp = man + (syu_local ** 2)
		if tmp > ans:
			ans = tmp
		
	print(ans)
	
	
if __name__ == "__main__":
	main()


Fixed submission:
# -*- coding: utf-8 -*-
import sys
import math
from bisect import bisect_left
from bisect import bisect_right
from collections import defaultdict
from heapq import heappop, heappush
import itertools
from decimal import *

input = sys.stdin.readline

def inputInt(): return int(input())
def inputMap(): return map(int, input().split())
def inputList(): return list(map(int, input().split()))
def inputStr(): return input()[:-1]

inf = float('inf')
mod = 1000000007

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-

def main():
	N,K = inputMap()
	sushi_mii = []
	sushi_kii = []
	syu = [[] for i in range(N)]
	for i in range(N):
		t,d = inputMap()
		t -= 1
		syu[t].append(d)

	for i in syu:
		if len(i) > 0:
			i.sort()
			i = i[::-1]
			for j,val in enumerate(i):
				if j == 0:
					sushi_mii.append(val)
				else:
					sushi_kii.append(val)

	sushi_mii.sort()
	sushi_mii = sushi_mii[::-1]
	sushi_kii.sort()
	sushi_kii = sushi_kii[::-1]

	sushi_mi = []
	sushi_ki = []
	for i,val in enumerate(sushi_mii):
		if i == 0:
			sushi_mi.append(val)
		else:
			tmp = sushi_mi[i-1]
			sushi_mi.append(tmp + val)

	for i,val in enumerate(sushi_kii):
		if i == 0:
			sushi_ki.append(val)
		else:
			tmp = sushi_ki[i-1]
			sushi_ki.append(tmp + val)

	ans = 0
	for i in range(1, min(len(sushi_mi), K)+1):
		x = 0
		tmp = 0
		if K - i > len(sushi_ki):
			continue

		tmp += sushi_mi[i-1]
		x = i

		if K - i > 0:
			tmp += sushi_ki[K - i-1]

		oishi = tmp + x * x
		if ans < oishi:
			ans = oishi
	print(ans)

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
if __name__ == "__main__":
	main()