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

mod = 10 ** 9 + 7

N, K = map(int, input().split())

dp = [[[0 for _ in range(2510)] for _ in range(51)] for _ in range(51)]
dp[0][0][0] = 1

for i in range(N):
    for j in range(N):
        for k in range(K + 1):
            dp[i + 1][j + 1][k + 2 * j + 2] += dp[i][j][k]
            if j > 0:
                dp[i + 1][j][k + 2 * j] += dp[i][j][k] * j
                dp[i + 1][j][k + 2 * j] += dp[i][j][k] * j

            dp[i + 1][j][k + 2 * j] += dp[i][j][k]

            if j > 0:
                dp[i + 1][j - 1][k + 2 * j - 2] += dp[i][j][k] * j * j


print(dp[N][0][K] % mod)

Fixed submission:
mod = 10 ** 9 + 7

N, K = map(int, input().split())

dp = [[[0 for _ in range(2610)] for _ in range(51)] for _ in range(51)]
dp[0][0][0] = 1

for i in range(N):
    for j in range(N):
        for k in range(K + 1):
            dp[i + 1][j + 1][k + 2 * j + 2] += dp[i][j][k]
            if j > 0:
                dp[i + 1][j][k + 2 * j] += dp[i][j][k] * j
                dp[i + 1][j][k + 2 * j] += dp[i][j][k] * j

            dp[i + 1][j][k + 2 * j] += dp[i][j][k]

            if j > 0:
                dp[i + 1][j - 1][k + 2 * j - 2] += dp[i][j][k] * j * j


print(dp[N][0][K] % mod)

Buggy python submission (verdict: Runtime Error):
mod = 10 ** 9 + 7

N, K = map(int, input().split())

dp = [[[0 for _ in range(2510)] for _ in range(51)] for _ in range(51)]
dp[0][0][0] = 1

for i in range(N):
    for j in range(N):
        for k in range(K + 1):
            dp[i + 1][j + 1][k + 2 * j + 2] += dp[i][j][k]
            if j > 0:
                dp[i + 1][j][k + 2 * j] += dp[i][j][k] * j
                dp[i + 1][j][k + 2 * j] += dp[i][j][k] * j

            dp[i + 1][j][k + 2 * j] += dp[i][j][k]

            if j > 0:
                dp[i + 1][j - 1][k + 2 * j - 2] += dp[i][j][k] * j * j


print(dp[N][0][K] % mod)

Fixed submission:
mod = 10 ** 9 + 7

N, K = map(int, input().split())

dp = [[[0 for _ in range(2610)] for _ in range(51)] for _ in range(51)]
dp[0][0][0] = 1

for i in range(N):
    for j in range(N):
        for k in range(K + 1):
            dp[i + 1][j + 1][k + 2 * j + 2] += dp[i][j][k]
            if j > 0:
                dp[i + 1][j][k + 2 * j] += dp[i][j][k] * j
                dp[i + 1][j][k + 2 * j] += dp[i][j][k] * j

            dp[i + 1][j][k + 2 * j] += dp[i][j][k]

            if j > 0:
                dp[i + 1][j - 1][k + 2 * j - 2] += dp[i][j][k] * j * j


print(dp[N][0][K] % mod)
Buggy python submission (verdict: Runtime Error):
from collections import Counter

class UnionFind:
    def __init__(self, size):
        self.data = [-1] * size

    def find(self, x):
        if self.data[x] < 0:
            return x
        else:
            self.data[x] = self.find(self.data[x])
            return self.data[x]

    def union(self, x, y):
        x, y = self.find(x), self.find(y)
        if x != y:
            if self.data[y] < self.data[x]:
                x, y = y, x
            self.data[x] += self.data[y]
            self.data[y] = x
        return (x != y)

    def same(self, x, y):
        return (self.find(x) == self.find(y))

    def size(self, x):
        return -self.data[self.find(x)]


MAX = 10 ** 5
uf = UnionFind(2 * MAX)

N, *XY = map(int, open(0).read().split())

for x, y in zip(*[iter(XY)] * 2):
    uf.union(x, y + MAX)

X = Counter(uf.find(i) for i in range(MAX))
Y = Counter(uf.find(i) for i in range(MAX, MAX * 2))
res = sum(X[i] * Y[i] for i in range(MAX * 2))

print(res - N)


Fixed submission:
from collections import Counter

class UnionFind:
    def __init__(self, size):
        self.data = [-1] * size

    def find(self, x):
        if self.data[x] < 0:
            return x
        else:
            self.data[x] = self.find(self.data[x])
            return self.data[x]

    def union(self, x, y):
        x, y = self.find(x), self.find(y)
        if x != y:
            if self.data[y] < self.data[x]:
                x, y = y, x
            self.data[x] += self.data[y]
            self.data[y] = x
        return (x != y)

    def same(self, x, y):
        return (self.find(x) == self.find(y))

    def size(self, x):
        return -self.data[self.find(x)]


MAX = 10 ** 5 + 10
uf = UnionFind(2 * MAX)

N, *XY = map(int, open(0).read().split())

for x, y in zip(*[iter(XY)] * 2):
    uf.union(x, y + MAX)

X = Counter(uf.find(i) for i in range(MAX))
Y = Counter(uf.find(i) for i in range(MAX, MAX * 2))
res = sum(X[i] * Y[i] for i in range(MAX * 2))

print(res - N)
Buggy python submission (verdict: Time Limit Exceeded):
def main():
    mod = 10 ** 9 + 7

    N, M = map(int, input().split())
    S = input().split()
    T = input().split()

    dp = [[0] * (M + 1) for _ in range(N + 1)]
    dp[0][0] = 1

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

    print(dp[N][M])


if __name__ == '__main__':
    main()


Fixed submission:
mod = 10 ** 9 + 7

N, M = map(int, input().split())
S = input().split()
T = input().split()

dp = [[0] * (M + 1) for _ in range(N + 1)]
dp[0][0] = 1

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

print(dp[N][M])

Buggy python submission (verdict: Time Limit Exceeded):
mod = 10 ** 9 + 7

N, M = map(int, input().split())
S = input().split()
T = input().split()

dp = [[0] * (M + 1) for _ in range(N + 1)]
dp[0][0] = 1

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

print(dp[N][M])


Fixed submission:
mod = 10 ** 9 + 7

N, M = map(int, input().split())
S = input().split()
T = input().split()

dp = [[0] * (M + 1) for _ in range(N + 1)]
dp[0][0] = 1

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

print(dp[N][M])
Buggy python submission (verdict: Runtime Error):
N = int(input())
R = [[int(i) for i in input().split()] for _ in range(N)]

from itertools import combinations
L = []
for (x1, y1), (x2, y2) in combinations(R, 2):
    L.append((x1 - x2, y1 - y2))
    L.append((-x1 + x2, -y1 + y2))

from collections import Counter
C = Counter(L)
print(N - C.most_common(1)[0][1])


Fixed submission:
from collections import Counter
from itertools import combinations

N = int(input())
R = [[int(i) for i in input().split()] for _ in range(N)]

if N == 1:
    print(1)
    quit()

L = []
for (x1, y1), (x2, y2) in combinations(R, 2):
    L.append((x1 - x2, y1 - y2))
    L.append((-x1 + x2, -y1 + y2))

C = Counter(L)
print(N - C.most_common(1)[0][1])
Buggy python submission (verdict: Runtime Error):
mod = 1000000007
 
N, M = [int(i) for i in input().split()]
A = [int(input()) for _ in range(M)] + [N + 1]
 
for i in range(len(A) - 1):
    if A[i] + 1 == A[i + 1]:
        print(0)
        quit()
 
memo = {0: 1, 1: 1, 2: 2, 3: 3}
 
def fib(n):
    if n not in memo:
        memo[n] = fib(n - 1) + fib(n - 2)
    return memo[n]
 
ans = 1
pre = -1
lo = 0
for a in A:
    f = fib(a - 1 - lo)
    ans = (ans * f) % mod
 
    lo = a + 1
    pre = a
 
print(ans)

Fixed submission:
mod = 1000000007

N, M = [int(i) for i in input().split()]
A = [int(input()) for _ in range(M)] + [N + 1]

for i in range(len(A) - 1):
    if A[i] + 1 == A[i + 1]:
        print(0)
        quit()

fib = [-1] * (N + 1)
fib[0] = 1
fib[1] = 1
for i in range(2, len(fib)):
    fib[i] = fib[i - 1] + fib[i - 2]


ans = 1
pre = -1
lo = 0
for a in A:
    f = fib[a - 1 - lo]
    ans = (ans * f) % mod

    lo = a + 1
    pre = a

print(ans)

Buggy python submission (verdict: Runtime Error):
mod = 1000000007

N, M = [int(i) for i in input().split()]
A = [int(input()) for _ in range(M)] + [N + 1]

for i in range(len(A) - 1):
    if A[i] + 1 == A[i + 1]:
        print(0)
        quit()

memo = {0: 1, 1: 1, 2: 2, 3: 3}

def fib(n):
    if n not in memo:
        memo[n] = fib(n - 1) + fib(n - 2)
    return memo[n]

ans = 1
pre = -1
lo = 0
for a in A:
    f = fib(a - 1 - lo)
    ans = (ans * f) % mod

    lo = a + 1
    pre = a

print(ans)

Fixed submission:
mod = 1000000007

N, M = [int(i) for i in input().split()]
A = [int(input()) for _ in range(M)] + [N + 1]

for i in range(len(A) - 1):
    if A[i] + 1 == A[i + 1]:
        print(0)
        quit()

fib = [-1] * (N + 1)
fib[0] = 1
fib[1] = 1
for i in range(2, len(fib)):
    fib[i] = fib[i - 1] + fib[i - 2]


ans = 1
pre = -1
lo = 0
for a in A:
    f = fib[a - 1 - lo]
    ans = (ans * f) % mod

    lo = a + 1
    pre = a

print(ans)
Buggy python submission (verdict: Time Limit Exceeded):
from bisect import bisect_left

H, W = [int(i) for i in input().split()]
S = [list(input()) for _ in range(H)]

R = [[] for i in range(H)]
for j in range(H):
    R[j].append(-1)
    for i in range(W):
        if S[j][i] == "#":
            R[j].append(i)
    R[j].append(W)

C = [[] for i in range(W)]
for j in range(W):
    C[j].append(-1)
    for i in range(H):
        if S[i][j] == "#":
            C[j].append(i)
    C[j].append(H)

ma = 0
for i in range(H):
    for j in range(W):
        if S[i][j] == "#":
            continue

        r = bisect_left(R[i], j)
        c = bisect_left(C[j], i)

        rr = R[i][r] - R[i][r - 1] - 1
        cc = C[j][c] - C[j][c - 1] - 1

        ma = max(ma, rr + cc - 1)

print(ma)

Fixed submission:
H, W = [int(i) for i in input().split()]
S = [input() for _ in range(H)]

L = [[0 for _ in range(W)] for _ in range(H)]
R = [[0 for _ in range(W)] for _ in range(H)]
U = [[0 for _ in range(W)] for _ in range(H)]
D = [[0 for _ in range(W)] for _ in range(H)]

for i in range(H):
    for j in range(W):
        if S[i][j] != "#":
            L[i][j] = 1 if j == 0 else 1 + L[i][j - 1]
    for j in range(W - 1, -1, -1):
        if S[i][j] != "#":
            R[i][j] = 1 if j == W - 1 else 1 + R[i][j + 1]

for j in range(W):
    for i in range(H):
        if S[i][j] != "#":
            U[i][j] = 1 if i == 0 else 1 + U[i - 1][j]
    for i in range(H - 1, -1, -1):
        if S[i][j] != "#":
            D[i][j] = 1 if i == H - 1 else 1 + D[i + 1][j]

ans = 0
for i in range(H):
    for j in range(W):
        ans = max(ans, L[i][j] + R[i][j] + U[i][j] + D[i][j] - 3)
print(ans)
Buggy python submission (verdict: Time Limit Exceeded):
mod = 10 ** 9 + 7

S = input()
N = len(S)

ans = 0
cnt = 0
length = N - 1
for s in S:
    if s == '1':
        ans += pow(2, cnt, mod) * pow(3, length)
        ans %= mod
        cnt += 1
    length -= 1
ans += pow(2, cnt, mod)
ans %= mod

print(ans)


Fixed submission:
mod = 10 ** 9 + 7

S = input()
N = len(S)

ans = 0
cnt = 0
length = N - 1
for s in S:
    if s == '1':
        ans += pow(2, cnt, mod) * pow(3, length, mod)
        ans %= mod
        cnt += 1
    length -= 1
ans += pow(2, cnt, mod)
ans %= mod

print(ans)

Buggy python submission (verdict: Time Limit Exceeded):
mod = 10 ** 9 + 7

S = input()
N = len(S)

ans = 0
cnt = 0
length = N - 1
for s in S:
    if s == '1':
        ans += pow(2, cnt, mod) * pow(3, length)
        ans %= mod
        cnt += 1
    length -= 1
ans += pow(2, cnt, mod)
ans %= mod

print(ans)


Fixed submission:
mod = 10 ** 9 + 7

S = input()
N = len(S)

ans = 0
cnt = 0
length = N - 1
for s in S:
    if s == '1':
        ans += pow(2, cnt, mod) * pow(3, length, mod)
        ans %= mod
        cnt += 1
    length -= 1
ans += pow(2, cnt, mod)
ans %= mod

print(ans)
Buggy python submission (verdict: Runtime Error):
from numpy.linalg import matrix_power

class mint:
    def __init__(self, n):
        self.n = n % mod

    def __add__(self, other):
        return mint(self.n + other.n)

    def __mul__(self, other):
        return mint(self.n * other.n)

    def __repr__(self):
        return str(self.n)


L, A, B, mod = map(int, input().split())

ans = [mint(0), mint(A), mint(B)]

size = 0
for d in range(1, 19):
    x = pow(10, d, mod)
    T = [
        [mint(x), mint(0), mint(0)],
        [mint(1), mint(1), mint(0)],
        [mint(0), mint(1), mint(1)],
    ]

    cnt = min(L - size, max(0, 10 ** d - 1 - A) // B - size + 1)
    if cnt > 0:
        ans = ans @ matrix_power(T, cnt)

    size += cnt
    if size >= L:
        break

print(ans[0])

Fixed submission:
from numpy.linalg import matrix_power
from numpy import dot

class mint:
    def __init__(self, n):
        self.n = n % mod

    def __add__(self, other):
        return mint(self.n + other.n)

    def __mul__(self, other):
        return mint(self.n * other.n)

    def __repr__(self):
        return str(self.n)


L, A, B, mod = map(int, input().split())

ans = [mint(0), mint(A), mint(B)]

size = 0
for d in range(1, 19):
    x = pow(10, d, mod)
    T = [
        [mint(x), mint(0), mint(0)],
        [mint(1), mint(1), mint(0)],
        [mint(0), mint(1), mint(1)],
    ]

    cnt = min(L - size, max(0, 10 ** d - 1 - A) // B - size + 1)
    if cnt > 0:
        ans = dot(ans, matrix_power(T, cnt))

    size += cnt
    if size >= L:
        break

print(ans[0])
Buggy python submission (verdict: Runtime Error):
N, A, B, C, D = [int(i) for i in input().split()]
S = input()

def solve():
    ac, bd = S[A - 1:C], S[B - 1:D]
    if A < B < D < C:
        bd = S[B - 2:D + 1]
        return ("..." in bd and "##" not in ac)
    else:
        return ("##" not in ac and "##" not in bd)

if solve():
    print("Yes")
else:
    print("No")

Fixed submission:
S = input().replace("BC", "D")

ans = a = 0
for s in S:
    if s == "A":
        a += 1
    elif s == "D":
        ans += a
    else:
        a = 0

print(ans)
Buggy python submission (verdict: Runtime Error):
N = int(input())
S, P = [""], [0]
for _ in range(N):
    s, p = input().split()
    S.append(s)
    P.append(int(p))

I = [*range(1, N + 1)]
I.sort(key=lambda i: P[i], reverse=True)
I.sort(key=lambda i: S[i])

print(*I, sep='\n')


Fixed submission:
N = int(input())

dic = {}
for i in range(N):
    S, P = [i for i in input().split()]
    if S in dic:
        dic[S].append((int(P), i + 1))
    else:
        dic[S] = [(int(P), i + 1)]

for s in sorted(dic.keys()):
    for l in reversed(sorted(dic[s])):
        print(l[1])

Buggy python submission (verdict: Runtime Error):
N = int(input())
S, P = [""], [0]
for _ in range(N):
    s, p = input().split()
    S.append(s)
    P.append(int(p))

I = [*range(1, N + 1)]
I.sort(key=lambda i: P[i], reverse=True)
I.sort(key=lambda i: S[i])

print(*I, sep='\n')


Fixed submission:
N = int(input())

dic = {}
for i in range(N):
    S, P = [i for i in input().split()]
    if S in dic:
        dic[S].append((int(P), i + 1))
    else:
        dic[S] = [(int(P), i + 1)]

for s in sorted(dic.keys()):
    for l in reversed(sorted(dic[s])):
        print(l[1])
Buggy python submission (verdict: Time Limit Exceeded):
R, G, B, N = [int(i) for i in input().split()]

count = 0
for r in range(N // R + 1):
    for g in range(N // G + 1):
        if (N - R * r - G * g) % B == 0 and (N - R * r - G * g) // B >= 0:
            count += 1

print(count)

Fixed submission:
R, G, B, N = [int(i) for i in input().split()]

count = 0
for r in range(N // R + 1):
    for g in range(N // G + 1):
        if (N - R * r - G * g) % B == 0 and (N - R * r - G * g) // B >= 0:
            count += 1

print(count)
Buggy python submission (verdict: Time Limit Exceeded):
N = int(input())
A = [int(i) for i in input().split()]

from fractions import gcd

L = [0, 0]
for ai in A:
    L.append(gcd(L[-1], ai))

R = [0, 0]
for ai in A[::-1]:
    R.insert(0, gcd(R[0], ai))

print(max(gcd(l, r) for l, r in zip(L, R)))


Fixed submission:
N = int(input())
A = [int(i) for i in input().split()]

from fractions import gcd

L = [0, 0]
for i in range(N):
    L.append(gcd(L[-1], A[i]))

R = [0, 0]
for i in range(N - 1, -1, -1):
    R.append(gcd(R[-1], A[i]))
R.reverse()

print(max(gcd(l, r) for l, r in zip(L, R)))
Buggy python submission (verdict: Runtime Error):
N = int(input())
A = [int(i) for i in input().split()]

B = map(abs, A)
if len([a for a in A if a < 0]) % 2 == 1:
    print(sum(B) - min(B))
else:
    print(sum(B))


Fixed submission:
N = int(input())
A = [int(i) for i in input().split()]

m = 0
su = 0
mi = 10 ** 10
for ai in A:
    if ai < 0:
        m += 1
    mi = min(mi, abs(ai))
    su += abs(ai)

if m % 2 == 1:
    print(su - 2 * mi)
else:
    print(su)

Buggy python submission (verdict: Runtime Error):
N = int(input())
A = [int(i) for i in input().split()]

A.sort()
su = 0
m = 0
for i in range(N):
    if A[i] < 0:
       su -= A[i]
       m += 1
    else:
       su += A[i]

if m % 2 == 1:
    print(su - 2 * A[m])
else:
    print(su)

Fixed submission:
N = int(input())
A = [int(i) for i in input().split()]

m = 0
su = 0
mi = 10 ** 10
for ai in A:
    if ai < 0:
        m += 1
    mi = min(mi, abs(ai))
    su += abs(ai)

if m % 2 == 1:
    print(su - 2 * mi)
else:
    print(su)
Buggy python submission (verdict: Time Limit Exceeded):
def comb(n):
    return n * (n - 1) // 2

class UnionFind():
    def __init__(self, size):
        self.table = [-1 for _ in range(size)]

    def find(self, x):
        if self.table[x] < 0:
            return x
        else:
            self.table[x] = self.find(self.table[x])
            return self.table[x]

    def union(self, x, y):
        s1 = self.find(x)
        s2 = self.find(y)
        if s1 != s2:
            if self.table[s1] <= self.table[s2]:
                self.table[s1] += self.table[s2]
                self.table[s2] = s1
            else:
                self.table[s2] += self.table[s1]
                self.table[s1] = s2
            return True
        return False

    def subsetall(self):
        a = []
        for i in range(len(self.table)):
            if self.table[i] < 0:
                a.append((i, -self.table[i]))
        return a

    def cost(self):
        result = 0
        for i in range(len(self.table)):
            if self.table[i] < 0:
                result += comb(-self.table[i])
        return result


N, M = [int(i) for i in input().split()]

A = [0] * M
B = [0] * M
for i in range(M):
    A[M - i - 1], B[M - i - 1] = [int(i) - 1 for i in input().split()]

uf = UnionFind(N)

result = comb(N)
lst = [result]
for i in range(M - 1):
    uf.union(A[i], B[i])
    lst.append(result - uf.cost())

lst.reverse()

for li in lst:
    print(li)


Fixed submission:
class UnionFind:
    def __init__(self, size):
        self.data = [-1 for _ in range(size)]
    def find(self, x):
        if self.data[x] < 0:
            return x
        else:
            self.data[x] = self.find(self.data[x])
            return self.data[x]
    def union(self, x, y):
        x, y = self.find(x), self.find(y)
        if x != y:
            if self.data[y] < self.data[x]:
                x, y = y, x
            self.data[x] += self.data[y]
            self.data[y] = x
        return (x != y)
    def same(self, x, y):
        return (self.find(x) == self.find(y))
    def size(self, x):
        return -self.data[self.find(x)]

N, M = [int(i) for i in input().split()]

A = [0] * M
B = [0] * M
for i in range(M):
    A[M - i - 1], B[M - i - 1] = [int(i) - 1 for i in input().split()]

uf = UnionFind(N)

result = N * (N - 1) // 2
lst = [result]
for i in range(M - 1):
    if not uf.same(A[i], B[i]):
        result -= uf.size(A[i]) * uf.size(B[i])
    lst.append(result)
    uf.union(A[i], B[i])

for i in range(M - 1, -1, -1):
    print(lst[i])
Buggy python submission (verdict: Runtime Error):
from functools import reduce
from functools import gcd

N = int(input())
A = [int(i) for i in input().split()]

print(reduce(gcd, A))


Fixed submission:
def gcd(a, b):
    if a < b:
        a, b = b, a
    if b == 0:
        return a
    c = a % b
    return gcd(b, c)

N = int(input())
A = [int(i) for i in input().split()]

result = 0
for ai in A:
    result = gcd(result, ai)

print(result)
Buggy python submission (verdict: Runtime Error):
from heapq import heappush, heappop

INF = float("inf")

N, K, *TD = map(int, open(0).read().split())

S = [[] for _ in range(N)]
for t, d in zip(*[iter(TD)] * 2):
    S[t - 1].append(d)

E = []
for s in S:
    if s:
        E.append(sorted(s, reverse=True))
    else:
        E.append([-INF])
E.sort(reverse=True)

cur = 0
Q = []
for i in range(K):
    cur += E[i][0]
    for j in range(1, len(E[i])):
        heappush(Q, -E[i][j])

for i in range(K, N):
    for j in range(len(E[i])):
        heappush(Q, -E[i][j])

res = cur + K * K
for x in range(K, 0, -1):
    v, w = E[x][0], -Q[0]
    if v < w:
        heappop(Q)
        cur += w
        heappush(Q, -v)
        cur -= v

    res = max(res, cur + x * x)

print(res)



Fixed submission:
from heapq import heappush, heappop, heapreplace

INF = 10 ** 9 + 5

N, K, *TD = map(int, open(0).read().split())

E = [[] for _ in range(N)]
for t, d in zip(*[iter(TD)] * 2):
    E[t - 1].append(d)

for e in E:
    if e:
        e.sort(reverse=True)
    else:
        e.append(-INF)
E.sort(key=lambda x: -x[0])

cur = 0
Q = []
for ei in E[:K]:
    cur += ei[0]
    for eij in ei[1:]:
        heappush(Q, -eij)

for ei in E[K:]:
    for eij in ei:
        heappush(Q, -eij)

res = cur + K * K
for x in reversed(range(1, K)):
    v, w = E[x][0], -Q[0]
    if v < w:
        heappop(Q)
        cur += w
        heappush(Q, -v)
        cur -= v

    res = max(res, cur + x * x)

print(res)
Buggy python submission (verdict: Runtime Error):
import sys
sys.setrecursionlimit(200000)

H, W = [int(_) for _ in input().split()]
S = [input() for _ in range(H)]

move = ((1, 0), (-1, 0), (0, 1), (0, -1))
visited = [[False] * W for _ in range(H)]
b, w = 0, 0

def dfs(i, j):
    global b, w

    visited[i][j] = True

    if S[i][j] == '#':
        b += 1
    else:
        w += 1

    for di, dj in move:
        ni, nj = i + di, j + dj
        if 0 <= ni < H and 0 <= nj < W:
            if S[i][j] != S[ni][nj] and not visited[ni][nj]:
                dfs(ni, nj)


result = 0
for i in range(H):
    for j in range(W):
        if visited[i][j]:
            continue
        b, w = 0, 0
        dfs(i, j)
        result += b * w
print(result)



Fixed submission:
N = int(input())
H = int(input())
W = int(input())

print((N - H + 1) * (N - W + 1))