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
# -*- coding:utf-8 -*-

import sys
input = sys.stdin.readline

def main():
  N, K = map(int, input().split())
  # sushis = sorted([tuple(map(int, input().split())) for _ in range(N)], key=lambda x:x[1], reverse=True)
  sushis = [tuple(map(int, input().split())) for _ in range(N)]
  sushis.sort(key=lambda x:x[1])
  curr = sushis[:K]
  leave = sushis[K:]
  base = lambda li: sum([l[1] for l in li])
  types_max = len(set([c[0] for c in sushis]))
  types = []
  dups = []

  for c in curr:
    if c[0] in types:
      dups += [c]
    else:
      types += [c[0]]
  types = set(types)

  ps = base(curr) + len(types)**2
  di = {}
  di[len(types)] = ps
  n_types = len(types) + 1
  for l in leave:
    if len(types) >= types_max:
      break
    if len(dups) == 0:
      break
    if not l[0] in types:
      dup = dups.pop()
      di[n_types] = ps - dup[1] + l[1] + n_types**2 - (n_types-1)**2
      ps = di[n_types]
      n_types += 1
      types.add(l[0])

  print(sorted(di.items(), key=lambda x:x[1])[-1][1])


if __name__=='__main__':
  main()

#!/usr/bin/env python3
# -*- coding:utf-8 -*-


if __name__=='__main__':
  N, K = map(int, input().split())
  sushis = sorted([list(map(int, input().split())) for _ in range(N)], key=lambda x:x[1])[::-1]

  curr = sushis[:K]
  leave = sushis[K:]
  base = lambda li: sum([l[1] for l in li])
  types = set([c[0] for c in curr])
  di = {}

  types_max = len(set([c[0] for c in sushis]))

  for i in range(len(types), K+1):
    di[str(i)] = base(curr) + i**2

    # next
    if i < types_max:
      tdi = {}
      remove_number = -1
      for j, t in enumerate(curr[::-1]):
        if not str(t[0]) in tdi.keys():
          tdi[str(t[0])] = -j-1
        else:
          remove_number = tdi[str(t[0])]
          break
      curr.pop(int(remove_number))

      for j, l in enumerate(leave):
        if not l[0] in types:
          types.union({l[0]})
          break
      leave = leave[:j]
      curr += [leave.pop(0)]
    else:
      break
  print(sorted(di.items(), key=lambda x:x[1])[-1][1])

#!/usr/bin/env python3
# -*- coding:utf-8 -*-

if __name__=='__main__':
  N, K = map(int, input().split())
  sushis = sorted([list(map(int, input().split())) for _ in range(N)], key=lambda x:x[1])[::-1]

  curr = sushis[:K]
  leave = sushis[K:]
  base = lambda li: sum([l[1] for l in li])
  types = set([c[0] for c in curr])
  di = {}

  types_max = len(set([c[0] for c in sushis]))

  for i in range(len(types), K+1):
    di[str(i)] = base(curr) + i**2

    # next
    if i < types_max:
      tdi = {}
      remove_number = -1
      for j, t in enumerate(curr[::-1]):
        if not str(t[0]) in tdi.keys():
          tdi[str(t[0])] = -j-1
        else:
          remove_number = tdi[str(t[0])]
          break
      curr.pop(int(remove_number))

      for j, l in enumerate(leave):
        if not l[0] in types:
          types.union({l[0]})
          break
      curr += [leave.pop(j)]
    else:
      break
    print(curr)
  print(di)
  print(sorted(di.items(), key=lambda x:x[1])[-1][1])

#!/usr/bin/env python3
# -*- coding:utf-8 -*-


if __name__=='__main__':
  N, K = map(int, input().split())
  sushis = sorted([list(map(int, input().split())) for _ in range(N)], key=lambda x:x[1])[::-1]

  curr = sushis[:K]
  leave = sushis[K:]
  base = lambda li: sum([l[1] for l in li])
  types = set([c[0] for c in curr])
  di = {}

  types_max = len(set([c[0] for c in sushis]))

  for i in range(len(types), K+1):
    di[str(i)] = base(curr) + i**2

    # next
    if i < types_max:
      tdi = {}
      remove_number = -1
      for j, t in enumerate(curr[::-1]):
        if not str(t[0]) in tdi.keys():
          tdi[str(t[0])] = -j-1
        else:
          remove_number = tdi[str(t[0])]
          break
      curr.pop(int(remove_number))

      for j, l in enumerate(leave):
        if not l[0] in types:
          types.union({l[0]})
          break
      leave[:j]
      curr += [leave.pop(0)]
    else:
      break
  print(sorted(di.items(), key=lambda x:x[1])[-1][1])

#!/usr/bin/env python3
# -*- coding:utf-8 -*-

import sys
input = sys.stdin.readline

def main():
  N, K = map(int, input().split())
  sushis = [tuple(map(int, input().split())) for _ in range(N)]
  sushis.sort(key=lambda x:x[1], reverse=True)
  curr = sushis[:K]
  leave = sushis[K:]
  base = lambda li: sum([l[1] for l in li])
  types_max = len(set([c[0] for c in sushis]))
  types = set()
  dups = []

  for c in curr:
    if c[0] in types:
      dups += [c]
    else:
      types.add(c[0])
  types = set(types)

  ps = base(curr) + len(types)**2
  di = {}
  di[len(types)] = ps
  n_types = len(types) + 1
  for l in leave:
    if len(types) >= types_max:
      break
    if len(dups) == 0:
      break
    if not l[0] in types:
      dup = dups.pop()
      di[n_types] = ps - dup[1] + l[1] + n_types**2 - (n_types-1)**2
      ps = di[n_types]
      n_types += 1
      types.add(l[0])

  print(sorted(di.items(), key=lambda x:x[1])[-1][1])


if __name__=='__main__':
  main()

Verdicts: Wrong Answer

Runtime Error

Wrong Answer

Wrong Answer

Accepted