Harryxun/testdata huggingface dataset:
Buggy python submission (verdict: Runtime Error):
str = input()

L, R = map(int, str.split(" "))

min = 2019

for i in range(L, R + 1):
    for j in range(i + 1, R + 1):
        k = i * j % 2019
        if min > k:
            min = k

print(min)

Fixed submission:
str = input()

N, A, B = str.split(" ")

a = int(N) * int(A)
b = int(B)

if a > b:
    print(b)
else:
    print(a)

Buggy python submission (verdict: Runtime Error):
str = input()

L, R = map(int, str.split(" "))

min = 2019

for i in range(L, R + 1):
    for j in range(i + 1, R + 1):
        k = i * j % 2019
        if min > k:
            min = k

print(min)

Fixed submission:
str = input()

N, A, B = str.split(" ")

a = int(N) * int(A)
b = int(B)

if a > b:
    print(b)
else:
    print(a)
Buggy python submission (verdict: Runtime Error):
###
### atcorder test program
###

### math class
class math:
    ### GCD
    def gcd(self, a, b):
        if b == 0:
            return a
        return self.gcd(b, a%b)

    def gcd_list(self, A):
        if not A:
            return 0
        for i in reversed(range(1, self.gcd(max(A),min(A))+1)):
            count = 0
            for j in A:
                if (j % i) == 0:
                    count += 1
            if count == len(A):
                return i
        return 0

### input num
#i = input()
#A, B, T = [int(x) for x in i.split()]
N = int(input())
A = [int(x) for x in input().split()]
#C = [int(x) for x in input().split()]

### input debug
#print("debug:{}".format(gcd(max(A),min(A))))
#print(list(reversed(range(1, gcd(max(A),min(A))+1))))

mathfunc = math()
anslist = []

for i in range(N):
    Rgcd = mathfunc.gcd_list(A[:i])
    Lgcd = mathfunc.gcd_list(A[i+1:])
    anslist.append(mathfunc.gcd(Rgcd,Lgcd))

### output
print(max(anslist))


Fixed submission:
## input num
i = input()
A, B, T = [int(x) for x in i.split()]

print((T//A)*B)
Buggy python submission (verdict: Runtime Error):
N = int(input())
V = [int(x) for x in input().split()]
C = [int(x) for x in input().split()]

XY = 0

for i in range(N):
    XY += np.maximum(0, V[i]-C[i])

print(XY)

Fixed submission:

N = int(input())
V = [int(x) for x in input().split()]
C = [int(x) for x in input().split()]

XY = 0

for i in range(N):
    if (V[i]-C[i]) > 0 :
        XY += V[i]-C[i]

print(XY)
Buggy python submission (verdict: Runtime Error):
x,a = map(int,input().split())
if x < a:
  print(0)
  else : 
    print(10)

Fixed submission:
x,a = map(int,input().split())
if x < a:
  print(0)
else : 
  print(10)

Buggy python submission (verdict: Runtime Error):
x,a = map(int,input().split())
    if x < a:
        print(0)
    else : 
        print(10)

Fixed submission:
x,a = map(int,input().split())
if x < a:
  print(0)
else : 
  print(10)
Buggy python submission (verdict: Runtime Error):
# coding:utf-8

line = input()
ls = line.split(" ")
n = ls[0]
a = ls[1]
b = ls[2]

print(b) if n * a >= b else print(n * a)

Fixed submission:
# coding:utf-8

line = input()
ls = line.split(" ")
n = int(ls[0])
a = int(ls[1])
b = int(ls[2])

print(b) if n * a >= b else print(n * a)
Buggy python submission (verdict: Runtime Error):
# coding:utf-8

n=input()
sum=0
for i in range(1, n):
    l_r=input()
    l=int(l_r.split(" ")[0])
    r=int(l_r.split(" ")[1])
    sum+=r-l+1


Fixed submission:
# coding:utf-8

n=int(input())
sum=0
for i in range(0, n):
    l_r=input()
    l=int(l_r.split(" ")[0])
    r=int(l_r.split(" ")[1])
    sum+=r-l+1

print(sum)
Buggy python submission (verdict: Runtime Error):
# coding;utf-8

n=int(input())
ans_dict={}
for i in range(0, n):
    a=input()
    if not a in ans_dict:
        ans_dict[a]=1
    else:
        ans_dict[a]+=1
sum=0
for thing in ans_dict:
    if thing%2:
        sum+=0


Fixed submission:
# coding;utf-8

n=int(input())
ans_dict={}
for i in range(0, n):
    a=input()
    if not a in ans_dict:
        ans_dict[a]=1
    else:
        ans_dict[a]+=1
sum=0
for thing in ans_dict:
    if ans_dict[thing]%2:
        sum+=1
print(sum)

Buggy python submission (verdict: Time Limit Exceeded):
# coding;utf-8

n=int(input())
ans_list=[]
for i in range(0, n):
    a=input()
    if a in ans_list:
        ans_list.remove(a)
    else:
        ans_list.append(a)
print(len(ans_list))


Fixed submission:
# coding;utf-8

n=int(input())
ans_dict={}
for i in range(0, n):
    a=input()
    if not a in ans_dict:
        ans_dict[a]=1
    else:
        ans_dict[a]+=1
sum=0
for thing in ans_dict:
    if ans_dict[thing]%2:
        sum+=1
print(sum)
Buggy python submission (verdict: Runtime Error):
n = int(input())
a = list(map(int, input().split()))
a.sort()
if a[int(n/2 + 1)] == a[int(n/2)]:
    answer = 0
else:
    answer = a[int(n/2)] - a[int(n/2 - 1)]
print(answer)


Fixed submission:
n = int(input())
a = list(map(int, input().split()))
a.sort()
answer = a[int(n/2)] - a[int(n/2 - 1)]
print(answer)
Buggy python submission (verdict: Runtime Error):
n, x = (int(c) for c in input().split())
l = (int(c) for c in input().split())
l.insert(0,0)
flag = 0
count = 0
for i in range(int(n)+1):
    count = count + int(l[i])

    if count > int(x):
        print(i)
        flag = 1
        break

if flag == 0:
    print(int(n)+1)

Fixed submission:
n, x = input().split()
l = input().split()
l.insert(0,0)
a = 0
flag = 0
for i in range(int(n)+1):
    a = a + int(l[i])

    if a > int(x):
        print(i)
        flag = 1
        break

if flag == 0:
    print(int(n)+1)
Buggy python submission (verdict: Runtime Error):
#!/usr/bin/env python3
import sys
from math import *
from itertools import *
from collections import *
from functools import *
from operator import *
try:
    from math import gcd
except Exception:
    from fractions import gcd

MOD = 1000000007  # type: int

def solve(N: int, K: int):
    ns = []
    for i in range(1, N):
        if N < i * i:
            break
        ns.append(i)
        ns.append(N // i)
    ns = sorted(list(set(ns)))

    m = {n: i for i, n in enumerate(ns)}
    l = [1] * len(ns)
    for i, (s, b) in enumerate(zip(ns, ns[1:])):
        l[i] = (N // s) - (N // b)
    l[len(ns) - 1] = 1

    dp = [[0] * len(ns) for _ in range(K+1)]
    dp[0][m[N]] = 1

    for i in range(K):
        for j in range(len(ns) - 2, -1, -1):
            dp[i][j] += dp[i][j + 1]
        for j, nn in enumerate(ns):
            nj = m[N // nn]
            dp[i+1][j] = (dp[i+1][j] + dp[i][nj] * l[j]) % MOD

    return sum(dp[-1]) % MOD


def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    K = int(next(tokens))  # type: int
    result = solve(N, K)
    print(result)

if __name__ == '__main__':
    main()

Fixed submission:
#!/usr/bin/env python3
import sys
from math import *
from itertools import *
from collections import *
from functools import *
from operator import *
try:
    from math import gcd
except Exception:
    from fractions import gcd

MOD = 1000000007  # type: int

def solve(N: int, K: int):
    ns = []
    for i in range(1, N + 1):
        if N < i * i:
            break
        ns.append(i)
        ns.append(N // i)
    ns = sorted(list(set(ns)))

    m = {n: i for i, n in enumerate(ns)}
    l = [1] * len(ns)
    for i, (s, b) in enumerate(zip(ns, ns[1:])):
        l[i] = (N // s) - (N // b)

    dp = [[0] * len(ns) for _ in range(K+1)]
    dp[0][m[N]] = 1

    for i in range(K):
        for j in range(len(ns) - 2, -1, -1):
            dp[i][j] += dp[i][j + 1]
        for j, nn in enumerate(ns):
            nj = m[N // nn]
            dp[i+1][j] = (dp[i+1][j] + dp[i][nj] * l[j]) % MOD

    return sum(dp[-1]) % MOD


def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    K = int(next(tokens))  # type: int
    result = solve(N, K)
    print(result)

if __name__ == '__main__':
    main()
Buggy python submission (verdict: Runtime Error):
#!/usr/bin/env python3
import sys
from collections.abc import Iterable
from math import *
from itertools import *
from collections import *
from functools import *
from operator import *
try:
    from math import gcd
except Exception:
    from fractions import gcd


def solve(N: int, K: int):
    if (N - 2) * (N - 1) < 2 * K:
        return -1

    g = [[1] * N for _ in range(N)]
    for i in range(N):
        g[i][i] = 0

    for i in range(K):
        g[i][i+1] = 0
        g[i+1][i] = 0

    ret = []
    for i in range(N):
        for j in range(N):
            if j <= i:
                continue

            if g[i][j] != 0:
                ret.append('{} {}'.format(i+1, j+1))

    ret = [str(len(ret))] + ret
    return '\n'.join(ret)


def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    K = int(next(tokens))  # type: int
    result = solve(N, K)
    if isinstance(result, Iterable) and not isinstance(result, str):
        result = '\n'.join([str(v) for v in result])
    print(result)

if __name__ == '__main__':
    main()


Fixed submission:
#!/usr/bin/env python3
import sys
from collections.abc import Iterable
from math import *
from itertools import *
from collections import *
from functools import *
from operator import *
try:
    from math import gcd
except Exception:
    from fractions import gcd


def solve(N: int, K: int):
    g = [[1] * N for _ in range(N)]
    for i in range(N):
        g[i][i] = 0

    k = K
    for i in range(N):
        if k == 0:
            break
        for j in range(N - 1):
            if j <= i:
                continue

            g[i][j] = 0
            g[j][i] = 0
            k -= 1
            if k == 0:
                break

    if 0 < k:
        return '-1'

    ret = []
    for i in range(N):
        for j in range(N):
            if j <= i:
                continue

            if g[i][j] != 0:
                ret.append('{} {}'.format(i+1, j+1))
    ret = [str(len(ret))] + ret
    return '\n'.join(ret)


def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    K = int(next(tokens))  # type: int
    result = solve(N, K)
    if isinstance(result, Iterable) and not isinstance(result, str):
        result = '\n'.join([str(v) for v in result])
    print(result)

if __name__ == '__main__':
    main()
Buggy python submission (verdict: Time Limit Exceeded):
#!/usr/bin/env python3
import sys
from collections.abc import Iterable
from math import *
from itertools import *
from collections import *
from functools import *
from operator import *
try:
    from math import gcd
except Exception:
    from fractions import gcd

MOD = 1000000007  # type: int

def solve(N: int, M: int, S: "List[int]", T: "List[int]"):
    dp = [[0] * (M + 1) for _ in range(N + 1)]

    for i in range(N):
        for j in range(M):
            dp[i + 1][j + 1] = (dp[i+1][j] + dp[i][j+1] - dp[i][j]) % MOD
            if S[i] == T[j]:
                dp[i+1][j+1] += dp[i][j] + 1
            dp[i+1][j+1] %= MOD
    return (dp[-1][-1] + 1) % MOD


def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    M = int(next(tokens))  # type: int
    S = [ int(next(tokens)) for _ in range(N) ]  # type: "List[int]"
    T = [ int(next(tokens)) for _ in range(M) ]  # type: "List[int]"
    result = solve(N, M, S, T)
    if isinstance(result, Iterable) and not isinstance(result, str):
        result = '\n'.join([str(v) for v in result])
    print(result)

if __name__ == '__main__':
    main()


Fixed submission:
#!/usr/bin/env python3
import sys
from math import *
from itertools import *
from collections import *
from functools import *
from operator import *
try:
    from math import gcd
except Exception:
    from fractions import gcd

MOD = 1000000007  # type: int

def solve(N: int, M: int, S: "List[int]", T: "List[int]"):
    dp = [[0] * (M + 1) for _ in range(N + 1)]

    for i in range(N):
        for j in range(M):
            dp[i + 1][j + 1] = (dp[i+1][j] + dp[i][j+1] - dp[i][j]) % MOD
            if S[i] == T[j]:
                dp[i+1][j+1] += dp[i][j] + 1
            dp[i+1][j+1] %= MOD
    return (dp[-1][-1] + 1) % MOD


def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    M = int(next(tokens))  # type: int
    S = [ int(next(tokens)) for _ in range(N) ]  # type: "List[int]"
    T = [ int(next(tokens)) for _ in range(M) ]  # type: "List[int]"
    result = solve(N, M, S, T)
    print(result)

if __name__ == '__main__':
    main()
Buggy python submission (verdict: Runtime Error):
import sys

import numpy as np

readInt = lambda: int(sys.stdin.readline())
readIntN = lambda: [int(v) for v in sys.stdin.readline().split(" ")]

total = 0
ret = 1
n = readInt()
aa = []
for i in range(n):
    v = readInt()
    aa.append(v)
    total += v
    ret = (ret * 3) % 998244353

memo0 = [[-1] * 100000 for _ in range(300)]
def dfs0(i, c):
    if 0 <= memo0[i][c]:
        return memo0[i][c]

    if i == len(aa):
        memo0[i][c] = 3 if total <= (2 * c) else 0
        return memo0[i][c]

    acc = max(0, memo0[i][c])
    acc = (acc + dfs0(i + 1, c)) % 998244353
    acc = (acc + dfs0(i + 1, c)) % 998244353
    acc = (acc + dfs0(i + 1, c + aa[i])) % 998244353
    memo0[i][c] = acc
    return memo0[i][c]

memo1 = [[-1] * 100000 for _ in range(300)]
def dfs1(i, c):
    if 0 <= memo1[i][c]:
        return memo1[i][c]

    if i == len(aa):
        memo1[i][c] = 3 if total == (2 * c) else 0
        return memo1[i][c]

    acc = max(0, memo1[i][c])
    acc = (acc + dfs1(i + 1, c)) % 998244353
    acc = (acc + dfs1(i + 1, c + aa[i])) % 998244353
    return memo1[i][c]

ret = (ret - dfs0(0, 0)) % 998244353
ret = (ret + dfs1(0, 0)) % 998244353
print(ret)

Fixed submission:
import sys
import numpy as np

readInt = lambda: int(sys.stdin.readline())
readIntN = lambda: [int(v) for v in sys.stdin.readline().split(' ')]

a, b, c = readIntN()
minv = min(a, b)
maxv = max(a, b)
ret = 'Yes' if minv < c < maxv else 'No'
print(ret)

Buggy python submission (verdict: Runtime Error):
import sys

import numpy as np

readInt = lambda: int(sys.stdin.readline())
readIntN = lambda: [int(v) for v in sys.stdin.readline().split(" ")]

total = 0
n = readInt()
aa = []
for i in range(n):
    v = readInt()
    aa.append(v)
    total += v

memo = {}
def dfs(a, b, c, i):
    if i == len(aa):
        ret = 1 if c < a + b else 0
        memo[(a, b, c)] = ret
        return ret
    if total < (2 * c):
        memo[(a, b, c)] = 0
        return 0
    if (a, b, c) in memo:
        return memo[(a, b, c)]

    ret = 0
    v0 = sorted([a + aa[i], b, c])
    ret = (ret + dfs(v0[0], v0[1], v0[2], i + 1)) % 998244353

    v1 = sorted([a, b + aa[i], c])
    ret = (ret + dfs(v1[0], v1[1], v1[2], i + 1)) % 998244353

    v2 = sorted([a, b, c + aa[i]])
    ret = (ret + dfs(v2[0], v2[1], v2[2], i + 1)) % 998244353

    memo[(a, b, c)] = ret
    return ret

print(dfs(0, 0, 0, 0))


Fixed submission:
import sys
import numpy as np

readInt = lambda: int(sys.stdin.readline())
readIntN = lambda: [int(v) for v in sys.stdin.readline().split(' ')]

a, b, c = readIntN()
minv = min(a, b)
maxv = max(a, b)
ret = 'Yes' if minv < c < maxv else 'No'
print(ret)
Buggy python submission (verdict: Runtime Error):
import sys
import numpy as np
import heapq

reatInt = lambda: int(sys.stdin.readline())
reatIntN = lambda: [int(v) for v in sys.stdin.readline().split(' ')]


n, k = reatIntN()
ss = sys.stdin.readline().strip()


class Node:
    def __init__(self, p, n, v):
        self.prev = p
        self.next = n
        self.value = v


r = Node(None, None, None)
r.next = r
r.prev = r

def append(n, value):
    nn = Node(n, n.next, value)
    n.next.prev = nn
    n.next = nn

def remove(n):
    pn = n.prev
    nn = n.next
    nn.prev = pn
    pn.next = nn
    n.next = None
    n.prev = None
    n.value = None


i = 0
while i < len(ss):
    b = i
    c = ss[i]
    while (i < len(ss)) and (c == ss[i]):
        i += 1
    append(r, (int(c), b, i))


def _calc(p):
    d = 0
    if (p is not r) and (p.value[0] == 0):
        lh = p.value[1]
        rh = p.value[2]
        if p.next is not r:
            lh = min(lh, p.next.value[1])
        if p.prev is not r:
            rh = max(rh, p.prev.value[2])
        d = rh - lh
    return d
 


p = r.next
pq = []
while p is not r:
    d = _calc(p)
    if 0 < d:
        heapq.heappush(pq, (-d, p.value[2], p))
    p = p.next


for i in range(k):
    if len(pq) == 0:
        break
    _, _, p = heapq.heappop(pq)
    if (p.value is not None) and (p.value[0] == 0):
        lh = p.value[1]
        rh = p.value[2]

        if p.next is not r:
            lh = min(lh, p.next.value[1])
            remove(p.next)
        if p.prev is not r:
            rh = max(rh, p.prev.value[2])
            remove(p.prev)
        p.value = (1, lh, rh)

        dp = _calc(p.prev)
        if 0 <= dp:
            heapq.heappush(pq, (-dp, p.prev.value[2], p.prev))
        dn = _calc(p.next)
        if 0 <= dn:
            heapq.heappush(pq, (-dn, p.next.value[2], p.next))

ret = 0
p = r.next
while p is not r:
    if p.value[0] == 1:
        ret = max(ret, p.value[2] - p.value[1])
    p = p.next
print(ret)

Fixed submission:
import sys
import numpy as np

reatInt = lambda: int(sys.stdin.readline())
reatIntN = lambda: [int(v) for v in sys.stdin.readline().split(' ')]


n, k = reatIntN()
ss = sys.stdin.readline().strip()

l = []
i = 0
if ss[0] == '0':
    l.append((1, 0))
else:
    l.append((0, 0))
while i < len(ss):
    c = ss[i]
    while (i < len(ss)) and (c == ss[i]):
        i += 1
    l.append((int(c), i))

i = 1
ret = 0
while i < len(l):
    q = 1 if  l[i][0] == 0 else 0
    b = i - 1
    e = min(i + 2 * k - q, len(l) - 1)
    ret = max(ret, l[e][1] - l[b][1])
    i += 1
print(ret)
Buggy python submission (verdict: Runtime Error):
import sys

import numpy as np

readInt = lambda: int(sys.stdin.readline())
readIntN = lambda: [int(v) for v in sys.stdin.readline().split(" ")]


n = readInt()
M = 10 ** 9 + 7


dp = np.zeros((n + 1, 256), dtype=np.uint32)
dp[0, 255] = 1


for i in range(n):
    for j4 in [0, 64, 128, 192]:
        for j3 in [0, 16, 32, 48]:
            j43 = j4 + j3
            for j2 in [0, 4, 8, 12]:
                j432 = j43 + j2
                for j1 in range(4):
                    j = j432 + j1
                    l321 = 4 * (j3 + j2 + j1)
                    for l0 in range(4):
                        if (6 == (j3 + 4 * j1 + l0)) or (6 == (j3 + j2 + l0)) or (6 == (4 * j2 + 4 * j1 + l0)) or (6 == (16 * j1 + j2 + l0)) or (6 == (4 * j2 + 4 * l0 + j1)):
                            continue
                        l = l321  + l0
                        dp[i + 1, l] = (dp[i + 1, l] + dp[i, j]) % M

ret = 0
for v in dp[-1]:
    ret = (ret + v) % M
print(ret)

Fixed submission:
import sys

b = sys.stdin.readline().strip()
r = {"A": "T", "T": "A", "C": "G", "G": "C"}[b]
print(r)
Buggy python submission (verdict: Time Limit Exceeded):
import sys

import numpy as np

readInt = lambda: int(sys.stdin.readline())
readIntN = lambda: [int(v) for v in sys.stdin.readline().split(" ")]


n = readInt()
M = 10 ** 9 + 7


dp = np.zeros((n + 1, 256), dtype=np.uint32)
dp[0, 255] = 1


for i in range(n):
    for l4 in range(4):
        for l3 in range(4):
            for l2 in range(4):
                for l1 in range(4):
                    j = 64 * l4 + 16 * l3 + 4 * l2 + l1
                    for l0 in range(4):
                        l = (64 * l3 + 16 * l2 + 4 * l1 + l0),

                        if 6 not in [
                            (16 * l2 + 4 * l1 + l0),
                            (16 * l3 + 4 * l1 + l0),
                            (16 * l1 + 4 * l2 + l0),
                            (16 * l2 + 4 * l0 + l1),
                            (16 * l3 + 4 * l2 + l0)
                        ]:
                            dp[i + 1, l] = (dp[i + 1, l] + dp[i, j]) % M

ret = 0
for v in dp[-1]:
    ret = (ret + v) % M
print(ret)

Fixed submission:
import sys

import numpy as np

readInt = lambda: int(sys.stdin.readline())
readIntN = lambda: [int(v) for v in sys.stdin.readline().split(" ")]


n = readInt()
M = 10 ** 9 + 7


dp = np.zeros((n + 1, 256), dtype=np.uint32)
dp[0, 255] = 1


for i in range(n):
    for j4 in [0, 64, 128, 192]:
        for j3 in [0, 16, 32, 48]:
            j43 = j4 + j3
            for j2 in [0, 4, 8, 12]:
                j432 = j43 + j2
                for j1 in range(4):
                    j = j432 + j1
                    l321 = 4 * (j3 + j2 + j1)
                    for l0 in range(4):
                        if (6 == (j3 + 4 * j1 + l0)) or (6 == (j3 + j2 + l0)) or (6 == (4 * j2 + 4 * j1 + l0)) or (6 == (16 * j1 + j2 + l0)) or (6 == (4 * j2 + 4 * l0 + j1)):
                            continue
                        l = l321  + l0
                        dp[i + 1, l] = (dp[i + 1, l] + dp[i, j]) % M

ret = 0
for v in dp[-1]:
    ret = (ret + v) % M
print(ret)
Buggy python submission (verdict: Runtime Error):
#!/usr/bin/env python3
import sys
from math import *
from itertools import *
from collections import *
from functools import *
from heapq import *
try:
    from math import gcd
except Exception:
    from fractions import gcd


def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    dt = sorted(zip(d, t), reverse=True)

    q0 = []

    acc = 0
    ts = 0
    p = [-1] * (N+1)
    for dd, tt in dt[:K]:
        acc += dd
        v = min(p[tt], dd)
        p[tt] = max(p[tt], dd)
        if 0 < v:
            heappush(q0, (v, tt))
        else:
            ts += 1
    
    q1 = []
    for dd, tt in dt[K:]:
        heappush(q1, (-dd, tt))

    mmm = acc + ts * ts
    while 0 < len(q1):
        d1, t1 = heappop(q1)
        if 0 < p[t1]:
            continue
        acc -= d1
        p[t1] = -d1
        ts += 1

        d0, _ = heappop(q0)
        acc -= d0
        mmm = max(mmm, acc + ts * ts)

    return mmm


def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    K = int(next(tokens))  # type: int
    t = [int()] * (N)  # type: "List[int]" 
    d = [int()] * (N)  # type: "List[int]" 
    for i in range(N):
        t[i] = int(next(tokens))
        d[i] = int(next(tokens))
    result = solve(N, K, t, d)
    print(result)

if __name__ == '__main__':
    main()


Fixed submission:
#!/usr/bin/env python3
import sys
from math import *
from itertools import *
from collections import *
from functools import *
from heapq import *
try:
    from math import gcd
except Exception:
    from fractions import gcd


def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    dt = sorted(zip(d, t), reverse=True)

    q0 = []

    acc = 0
    ts = 0
    p = [-1] * (N+1)
    for dd, tt in dt[:K]:
        acc += dd
        v = min(p[tt], dd)
        p[tt] = max(p[tt], dd)
        if 0 < v:
            heappush(q0, (v, tt))
        else:
            ts += 1

    q1 = []
    for dd, tt in dt[K:]:
        heappush(q1, (-dd, tt))

    mmm = acc + ts * ts
    while 0 < len(q1) and 0 < len(q0):
        d1, t1 = heappop(q1)
        if 0 < p[t1]:
            continue
        acc -= d1
        p[t1] = -d1
        ts += 1

        d0, _ = heappop(q0)
        acc -= d0
        mmm = max(mmm, acc + ts * ts)

    return mmm


def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    K = int(next(tokens))  # type: int
    t = [int()] * (N)  # type: "List[int]" 
    d = [int()] * (N)  # type: "List[int]" 
    for i in range(N):
        t[i] = int(next(tokens))
        d[i] = int(next(tokens))
    result = solve(N, K, t, d)
    print(result)

if __name__ == '__main__':
    main()