Harryxun/testdata huggingface dataset:
Buggy python submission (verdict: Time Limit Exceeded):
[N,X,Y] = [int(s) for s in input().rstrip().split(' ')]
L = [] #L[i][j]はi,j間の最短距離(i<=j)
for i in range(N):
  L.append([10**4 for j in range(N-i)])
for i in range(N):
  L[i][0] = 0
  if i < N-1:
    L[i][1] = 1
L[X-1][Y-X] = 1
for k in range(1,N):
  a = 0
  for i in range(N):
    for j in range(N-i):
      if L[i][j] == k:
        a += 1
      elif L[i][j] > k:
        for p in range(N):
          if p < i:
            if L[p][i-p] + L[p][i+j-p] < L[i][j]:
              L[i][j] = L[p][i-p] + L[p][i+j-p]
          elif p < i+j:
            if L[i][p-i] + L[p][i+j-p] < L[i][j]:
              L[i][j] = L[i][p-i] + L[p][i+j-p]
          else:
            if L[i][p-i] + L[i+j][p-i-j] < L[i][j]:
              L[i][j] = L[i][p-i] + L[i+j][p-i-j]
  print(a)


Fixed submission:
[N,X,Y] = [int(s) for s in input().rstrip().split(' ')]
L = [0 for n in range(N)]
for i in range(N):
  for j in range(N-i):
    L[min(j,abs(X-1-i)+abs(Y-1-i-j)+1,abs(X-1-j-i)+abs(Y-1-i)+1)]+=1
for k in range(1,N):
  print(L[k])

Buggy python submission (verdict: Time Limit Exceeded):
[N,X,Y] = [int(s) for s in input().rstrip().split(' ')]
L = [] #L[i][j]はi,j間の最短距離(i<=j)
for i in range(N):
  L.append([10**4 for j in range(N-i)])
for i in range(N):
  L[i][0] = 0
  if i < N-1:
    L[i][1] = 1
L[X-1][Y-X] = 1
for k in range(1,N):
  a = 0
  for i in range(N):
    for j in range(N-i):
      if L[i][j] == k:
        a += 1
      elif L[i][j] > k:
        for p in range(N):
          if p < i:
              L[i][j] = min(L[i][j],L[p][i-p] + L[p][i+j-p])
          elif p < i+j:
              L[i][j] = min(L[i][j],L[i][p-i] + L[p][i+j-p])
          else:
              L[i][j] = min(L[i][j],L[i][p-i] + L[i+j][p-i-j])
  print(a)

Fixed submission:
[N,X,Y] = [int(s) for s in input().rstrip().split(' ')]
L = [0 for n in range(N)]
for i in range(N):
  for j in range(N-i):
    L[min(j,abs(X-1-i)+abs(Y-1-i-j)+1,abs(X-1-j-i)+abs(Y-1-i)+1)]+=1
for k in range(1,N):
  print(L[k])

Buggy python submission (verdict: Time Limit Exceeded):
[N,X,Y] = [int(s) for s in input().rstrip().split(' ')]
L = [] #L[i][j]はi,j間の最短距離(i<=j)
for i in range(N):
  L.append([10**4 for j in range(N-i)])
for i in range(N):
  L[i][0] = 0
  if i < N-1:
    L[i][1] = 1
L[X-1][Y-X] = 1
for k in range(1,N):
  a = 0
  for i in range(N):
    for j in range(N-i):
      if L[i][j] == k:
        a += 1
      elif L[i][j] > k:
        for p in range(N):
          if p < i:
            if L[p][i-p] + L[p][i+j-p] < L[i][j]:
              L[i][j] = L[p][i-p] + L[p][i+j-p]
          elif p < i+j:
            if L[i][p-i] + L[p][i+j-p] < L[i][j]:
              L[i][j] = L[i][p-i] + L[p][i+j-p]
          else:
            if L[i][p-i] + L[i+j][p-i-j] < L[i][j]:
              L[i][j] = L[i][p-i] + L[i+j][p-i-j]
  print(a)

Fixed submission:
[N,X,Y] = [int(s) for s in input().rstrip().split(' ')]
L = [0 for n in range(N)]
for i in range(N):
  for j in range(N-i):
    L[min(j,abs(X-1-i)+abs(Y-1-i-j)+1,abs(X-1-j-i)+abs(Y-1-i)+1)]+=1
for k in range(1,N):
  print(L[k])
Buggy python submission (verdict: Time Limit Exceeded):
N = int(input())
D = {}
for i in range(N):
  a_i = input()
  if a_i in D.keys():
    D[a_i] += 1
  else:
    D[a_i] = 1
D_list = [kv[0] for kv in D.items() if kv[1] == max(D.values())]
A = sorted(D_list)
for x in A:
  print(x)

Fixed submission:
N = int(input())
D = {}
for i in range(N):
  a_i = input()
  if a_i in D.keys():
    D[a_i] += 1
  else:
    D[a_i] = 1
m = 0
for v in D.values():
	m = max(m, v)
ans = []
for k, v in D.items():
	if v == m:
		ans.append(k)
ans.sort()
print(*ans, sep="\n")

Buggy python submission (verdict: Time Limit Exceeded):
N = int(input())
D = {}
for i in range(N):
  a_i = input()
  if a_i in D.keys():
    D[a_i] += 1
  else:
    D[a_i] = 1
D_list = [kv[0] for kv in D.items() if kv[1] == max(D.values())]
A = sorted(D_list)
for x in A:
  print(x)

Fixed submission:
N = int(input())
D = {}
for i in range(N):
  a_i = input()
  if a_i in D.keys():
    D[a_i] += 1
  else:
    D[a_i] = 1
m = 0
for v in D.values():
	m = max(m, v)
ans = []
for k, v in D.items():
	if v == m:
		ans.append(k)
ans.sort()
print(*ans, sep="\n")
Buggy python submission (verdict: Time Limit Exceeded):
[N,K] = [int(s) for s in input().rstrip().split(' ')]
L_A = [int(s) for s in input().rstrip().split(' ')]
if N < K:
  print("error")
elif N == K:
  print((sum(L_A) + N) / 2)
else:
  a = N - K
  p = 0
  for i in range(a + 1):
    x = 0
    for j in range(K):
      x += L_A[i + j]
    p = max(p,x)
  print((p + K) / 2)

Fixed submission:
[N,K] = [int(s) for s in input().rstrip().split(' ')]
L_A = [int(s) for s in input().rstrip().split(' ')]
if N < K:
  print("error")
elif N == K:
  print((sum(L_A) + N) / 2)
else:
  a = N - K
  p = 0
  for i in range(K):
    p += L_A[i]
  q = p
  for j in range(N-K):
    q += L_A[j+K]-L_A[j]
    p = max(p,q)
  print((p + K) / 2)
Buggy python submission (verdict: Time Limit Exceeded):
[H,N] = [int(s) for s in input().rstrip().split(' ')] #HもんすたーのHP,N魔法の種類
L_1 = []
L_2 = []
for i in range(N):
    A_i = [int(s) for s in input().rstrip().split(' ')] #[攻撃力,消費魔力]
    L_1.append(A_i[0])
    L_2.append(A_i[1])
M = [10**8 for x in range(H+max(L_1))] #最小消費魔力
M[0] = 0
for k in range(1,H+max(L_1)):
  for n in range(N):
     M[k] = min(M[k],L_2[n]+M[max(k-L_1[n],0)])
print(min(M[H:]))

Fixed submission:
[H,N] = [int(s) for s in input().rstrip().split(' ')] #HもんすたーのHP,N魔法の種類
L_1 = []
L_2 = []
for i in range(N):
    A_i = [int(s) for s in input().rstrip().split(' ')] #[攻撃力,消費魔力]
    L_1.append(A_i[0])
    L_2.append(A_i[1])
M = [10**8 for x in range(H+max(L_1))] #最小消費魔力
M[0] = 0
for k in range(1,H+max(L_1)):
  for n in range(N):
     M[k] = min(M[k],L_2[n]+M[max(k-L_1[n],0)])
print(min(M[H:]))
Buggy python submission (verdict: Runtime Error):
[N,M] = [int(s) for s in input().rstrip().split(" ")
if N - M == 0:
	print("Yes")
else:
    print("No")

Fixed submission:
[N,M] = [int(s) for s in input().rstrip().split(" ")]
if N == M:
  print("Yes")
else:
  print("No")
Buggy python submission (verdict: Runtime Error):
[a,b] = [int(s) for s in input().rstrip().split(' ')]
W = ""
if a < b:
  for i in range(b):
    W.append(str(a))
  print(W)
else:
  for i in range(a):
    W.append(str(b))
  print(W) 
  

Fixed submission:
[a,b] = [int(s) for s in input().rstrip().split(' ')]
W = ""
if a < b:
  for i in range(b):
    W += str(a)
  print(W)
else:
  for i in range(a):
    W += str(b)
  print(W)
Buggy python submission (verdict: Time Limit Exceeded):
import sys
from heapq import heappush, heappop

inf = float('inf')

n, x, y = map(int, sys.stdin.readline().split())
g = [[] for _ in range(n)]
for i in range(n-1):
    g[i].append(i+1)
    g[i+1].append(i)
g[x-1].append(y-1)
g[y-1].append(x-1)

def dijkstra(i):
    dist = [inf] * n
    dist[i] = 0
    q = [(0, i)]
    visited = set()

    while q:
        d, x = heappop(q)
        if x in visited:
            continue
        dist[x] = d
        visited.add(x)
        for y in g[x]:
            if y in visited:
                continue
            heappush(q, (d+1, y))
    return dist
    
def main():
    res = [0] * (n + 1)
    for i in range(n):
        dist = dijkstra(i)
        for d in dist:
            res[d] += 1
    for c in res[1:-1]:
        print(c // 2)

if __name__ == "__main__":
    main()

Fixed submission:
import sys
from itertools import combinations

n, x, y = map(int, sys.stdin.readline().split())
    
def main():
    res = [0] * (n + 1)
    for i, j in combinations(range(1, n+1), 2):
        d = min(j - i, abs(i - x) + 1 + abs(j - y))
        res[d] += 1
        
    for i in range(1, n):
        print(res[i])

if __name__ == "__main__":
    main()

Buggy python submission (verdict: Time Limit Exceeded):
import sys
from collections import deque

inf = float('inf')

n, x, y = map(int, sys.stdin.readline().split())
g = [[] for _ in range(n)]
for i in range(n-1):
    g[i].append(i+1)
    g[i+1].append(i)
g[x-1].append(y-1)
g[y-1].append(x-1)

def bfs(i):
    dist = [inf] * n
    dist[i] = 0
    q = deque()
    q.append((i, 0))
    visited = set()

    while q:
        x, d = q.popleft()
        dist[x] = min(dist[x], d)
        visited.add(x)
        for y in g[x]:
            if y in visited:
                continue
            q.append((y, d+1))
    return dist
        

def main():
    res = [0] * (n + 1)
    for i in range(n):
        dist = bfs(i)
        for d in dist:
            res[d] += 1
    for c in res[1:-1]:
        print(c // 2)

if __name__ == "__main__":
    main()

Fixed submission:
import sys
from itertools import combinations

n, x, y = map(int, sys.stdin.readline().split())
    
def main():
    res = [0] * (n + 1)
    for i, j in combinations(range(1, n+1), 2):
        d = min(j - i, abs(i - x) + 1 + abs(j - y))
        res[d] += 1
        
    for i in range(1, n):
        print(res[i])

if __name__ == "__main__":
    main()
Buggy python submission (verdict: Runtime Error):
import sys
import numpy as np

a = np.array([sys.stdin.readline().split() for _ in range(3)], dtype=np.int64)
n = int(sys.stdin.readline().rstrip())
b = np.array(sys.stdin.read().split(), dtype=np.int64)
def main():
    c = np.isin(a, b)
    bl1 = np.any(np.all(c, axis=1))
    bl2 = np.any(np.all(c, axis=0))
    bl3 = np.all(np.diagonal(c))
    bl4 = np.all(np.diagonal(np.fliplr(c)))
    bl = bl1 | bl2 | bl3 | bl4
    print('Yes' if bl else 'No')

if __name__ ==  '__main__':
    main()

Fixed submission:
import sys

a = [list(map(int, sys.stdin.readline().split())) for _ in range(3)]
n, *b = map(int, sys.stdin.read().split())
b = set(b)

def main():
    res = [[False] * 3 for _ in range(3)]
    for i in range(3):
        for j in range(3):
            if a[i][j] in b:
                res[i][j] = True
    
    for i in range(3):
        for j in range(3):
            if not res[i][j]:
                break
        else:
            return 'Yes'
    
    for j in range(3):
        for i in range(3):
            if not res[i][j]:
                break
        else:
            return 'Yes'
    
    for i, j in zip(range(3), range(3)):
        if not res[i][j]:
            break
    else:
        return 'Yes'
    
    for i, j in zip(range(3), range(2, -1, -1)):
        if not res[i][j]:
            break
    else:
        return 'Yes'
    
    return 'No'

if __name__ == '__main__':
    ans = main()
    print(ans)
Buggy python submission (verdict: Runtime Error):
import sys
from collections import Counter

n, *s = sys.stdin.read().split()

def main():
    c = Counter(s)
    m = max(c.values())
    ans = [s for s, c in cnt.items() if c == m]
    print(*sorted(ans), sep='\n')
         
if __name__ ==  '__main__':
    main()

Fixed submission:
import sys
from collections import Counter

n = int(sys.stdin.readline().rstrip())
s = sys.stdin.read().split()

def main():
    c = Counter(s)
    ma = max(c.values())
    res = []
    for l, v in c.items():
        if v == ma:
            res.append(l)
    return sorted(res)

if __name__ == '__main__':
    ans = main()
    print(*ans, sep='\n')

Buggy python submission (verdict: Runtime Error):
import sys
from collections import Counter

n, *s = sys.stdin.read().split()

def main():
    c = Counter(s)
    m = max(c.values())
    ans = [s for s, c in cnt.items if c == m]
    print(*sorted(ans), sep='\n')
         
if __name__ ==  '__main__':
    main()

Fixed submission:
import sys
from collections import Counter

n = int(sys.stdin.readline().rstrip())
s = sys.stdin.read().split()

def main():
    c = Counter(s)
    ma = max(c.values())
    res = []
    for l, v in c.items():
        if v == ma:
            res.append(l)
    return sorted(res)

if __name__ == '__main__':
    ans = main()
    print(*ans, sep='\n')
Buggy python submission (verdict: Runtime Error):
import sys

MOD = 10 ** 9 + 7

def make_table(size=10**6+2, p=MOD):
    fac = [None] * (size + 1)
    fac[0] = 1
    for i in range(size):
        fac[i+1] = fac[i] * (i + 1) % p
    ifac = [None] * (size + 1)
    ifac[size] = pow(fac[size], p-2, p)
    for i in range(size, 0, -1):
        ifac[i-1] = ifac[i] * i % p
    return fac, ifac

fac, ifac = make_table()

def comb(n, r, mod=MOD):
    if r > n or r < 0:
        return 0
    return fac[n] * ifac[r] % mod * ifac[n-r] % mod

def g(y, x):
    return comb(y+x+2, y+1) - 1

r1, c1, r2, c2 = map(int, sys.stdin.readline().split())

def main():
    ans = (g(r2, c2) - (g(r1-1, c2) + g(r2, c1-1) - g(r1-1, c1-1))) % MOD
    return ans

if __name__ == '__main__':
    ans = main()
    print(ans)

Fixed submission:
import sys

MOD = 10 ** 9 + 7

def make_table(size=2*10**6+2, p=MOD):
    fac = [None] * (size + 1)
    fac[0] = 1
    for i in range(size):
        fac[i+1] = fac[i] * (i + 1) % p
    ifac = [None] * (size + 1)
    ifac[size] = pow(fac[size], p-2, p)
    for i in range(size, 0, -1):
        ifac[i-1] = ifac[i] * i % p
    return fac, ifac

fac, ifac = make_table()

def comb(n, r, mod=MOD):
    if r > n or r < 0:
        return 0
    return fac[n] * ifac[r] % mod * ifac[n-r] % mod

def g(y, x):
    return comb(y+x+2, y+1) - 1

r1, c1, r2, c2 = map(int, sys.stdin.readline().split())

def main():
    ans = (g(r2, c2) - (g(r1-1, c2) + g(r2, c1-1) - g(r1-1, c1-1))) % MOD
    return ans

if __name__ == '__main__':
    ans = main()
    print(ans)

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

MOD = 10 ** 9 + 7

def make_table(size=10**6+2, p=MOD):
    fac = [None] * (size + 1)
    fac[0] = 1
    for i in range(size):
        fac[i+1] = fac[i] * (i + 1) % p
    ifac = [None] * (size + 1)
    ifac[size] = pow(fac[size], p-2, p)
    for i in range(size, 0, -1):
        ifac[i-1] = ifac[i] * i % p
    return fac, ifac

fac, ifac = make_table()

def comb(n, r, mod=MOD):
    if r > n or r < 0:
        return 0
    return fac[n] * ifac[r] % mod * ifac[n-r] % mod

def g(y, x):
    return comb(y+x+2, y+1) - 1

r1, c1, r2, c2 = map(int, sys.stdin.readline().split())

def main():
    ans = g(r2, c2) - (g(r1-1, c2) + g(r2, c1-1) - g(r1-1, c1-1)) % MOD
    return ans

if __name__ == '__main__':
    ans = main()
    print(ans)

Fixed submission:
import sys

MOD = 10 ** 9 + 7

def make_table(size=2*10**6+2, p=MOD):
    fac = [None] * (size + 1)
    fac[0] = 1
    for i in range(size):
        fac[i+1] = fac[i] * (i + 1) % p
    ifac = [None] * (size + 1)
    ifac[size] = pow(fac[size], p-2, p)
    for i in range(size, 0, -1):
        ifac[i-1] = ifac[i] * i % p
    return fac, ifac

fac, ifac = make_table()

def comb(n, r, mod=MOD):
    if r > n or r < 0:
        return 0
    return fac[n] * ifac[r] % mod * ifac[n-r] % mod

def g(y, x):
    return comb(y+x+2, y+1) - 1

r1, c1, r2, c2 = map(int, sys.stdin.readline().split())

def main():
    ans = (g(r2, c2) - (g(r1-1, c2) + g(r2, c1-1) - g(r1-1, c1-1))) % MOD
    return ans

if __name__ == '__main__':
    ans = main()
    print(ans)
Buggy python submission (verdict: Time Limit Exceeded):
import sys

inf = float('inf')

h, n, *ab = map(int, sys.stdin.read().split())
ab = list(zip(*[iter(ab)] * 2))

def main():
    dp = [inf] * (h + 1)
    dp[0] = 0
    for i in range(1, h+1):
        for a, b in ab:
            dp[i] = min(dp[i], dp[max(i-a, 0)] + b)
        
    return dp[h]

if __name__ == '__main__':
    ans = main()
    print(ans)

Fixed submission:
import sys
import numpy as np

inf = float('inf')

h, n, *ab = map(int, sys.stdin.read().split())
a, b = np.array(ab, dtype=np.int64).reshape(n, 2).T

def main():
    dp = np.zeros(h+1, dtype=np.int64)
    for i in range(1, h+1):
        dp[i] = np.amin(dp[np.maximum(i-a, 0)] + b)
        
    return dp[h]

if __name__ == '__main__':
    ans = main()
    print(ans)
Buggy python submission (verdict: Time Limit Exceeded):
import sys
from functools import reduce

MOD = 10 ** 9 + 7

def gcd(a, b):
    while b:
        a, b = b, a % b
    return abs(a)

def lcm(a, b):
    return abs(a // gcd(a, b) * b)

n, *a = map(int, sys.stdin.read().split())

def main():
    l = reduce(lcm, a, 1)
    res = 0
    for i in a:
        res += l // i % MOD
        res %= MOD
    return res
    
if __name__ == '__main__':
    ans = main()
    print(ans)

Fixed submission:
import sys

MOD = 10 ** 9 + 7

from math import floor, sqrt
from bisect import bisect_right as bi_r
from collections import defaultdict

def prime_nums(n=10**6):
    sieve = set(range(2, n + 1))
    non_prime = set(range(2 * 2, n + 1, 2))
    sieve -= non_prime
    for i in range(3, floor(sqrt(n)) + 1, 2):
        if i in sieve:
            non_prime = set(range(i * 2, n + 1, i))
            sieve -= non_prime
    return sieve

p_nums = sorted(prime_nums())

def make_inv_table(n=10**6, p=MOD):
    fac = [None] * (n + 1)
    fac[0] = 1
    for i in range(n):
        fac[i+1] = fac[i] * (i + 1) % p
    ifac = [None] * (n + 1)
    ifac[n] = pow(fac[n], p-2, p)
    for i in range(n, 0, -1):
        ifac[i-1] = ifac[i] * i % p
    inv = [None] * (n + 1)
    for i in range(1, n+1):
        inv[i] = ifac[i] * fac[i-1] % p
    return inv

inv = make_inv_table()

n, *A = map(int, sys.stdin.read().split())

def main():
    facs = defaultdict(int)
    for a in A:
        if a == 1:
            continue
        for p in p_nums[:bi_r(p_nums, sqrt(a))]:
            cnt = 0
            while not a % p:
                cnt += 1
                a //= p
            facs[p] = max(facs[p], cnt)
            if a == 1:
                break
        else:
            facs[a] = max(facs[a], 1)
    
    lcm = 1
    for f, c in facs.items():
        lcm *= pow(f, c, MOD)
        lcm %= MOD
    
    ans = 0
    for a in A:
        ans += lcm * inv[a] % MOD
        ans %= MOD
    return ans

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

MOD = 10 ** 9 + 7

def make_table(size=10**6, p=MOD):
    fac = [None] * (size + 1)
    fac[0] = 1
    for i in range(size):
        fac[i+1] = fac[i] * (i + 1) % p
    ifac = [None] * (size + 1)
    ifac[size] = pow(fac[size], p-2, p)
    for i in range(size, 0, -1):
        ifac[i-1] = ifac[i] * i % p
    return fac, ifac

fac, ifac = make_table(10**5)

def comb(n, r, mod=MOD):
    if r > n or r < 0:
        return 0
    return fac[n] * ifac[r] % mod * ifac[n-r] % mod

n, *c = map(int, sys.stdin.read().split())

def main():
    c.sort()
    res = 0
    for i in range(n):
        cnt = 0
        for d in range(1, n+1):
            for j in range(d):
                cnt += comb(i, j) * comb(n-i-1, d-j-1) * (d - j) % MOD
        
        res += c[i] * cnt
        res %= MOD
    return res * 2 ** n % MOD

if __name__ == '__main__':
    ans = main()
    print(ans)

Fixed submission:
import sys

MOD = 10 ** 9 + 7

def make_table(size=10**6, p=MOD):
    fac = [None] * (size + 1)
    fac[0] = 1
    for i in range(size):
        fac[i+1] = fac[i] * (i + 1) % p
    ifac = [None] * (size + 1)
    ifac[size] = pow(fac[size], p-2, p)
    for i in range(size, 0, -1):
        ifac[i-1] = ifac[i] * i % p
    return fac, ifac

fac, ifac = make_table(2*10**5)

def comb(n, r, mod=MOD):
    if r > n or r < 0:
        return 0
    return fac[n] * ifac[r] % mod * ifac[n-r] % mod

n, *c = map(int, sys.stdin.read().split())

def main():
    c.sort(reverse=True)
    res = 0
    # for i in range(n):
    #     cnt = 0
    #     for d in range(1, n+1):
    #         for j in range(d):
    #             cnt += comb(i, j) * comb(n-i-1, d-j-1) * (d - j) % MOD
        
    #     res += c[i] * cnt
    #     res %= MOD
    for i in range(n):
        res += c[i] * (i + 2)
        res %= MOD
    return res * 4 ** (n - 1) % MOD

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

n, a, b = map(int, sys.stdin.readline().split())

def main():
    d = b - a
    if d % 2 == 0:
        return d // 2
    else:
        return min(n - b, abs(1 - a)) + 1 + (d - 1) // 2)

if __name__ == '__main__':
    ans = main()
    print(ans)

Fixed submission:
import sys

n, a, b = map(int, sys.stdin.readline().split())

def main(a, b):
    d = b - a
    if d % 2 == 0:
        return d // 2

    if n - b >= abs(1 - a):
        res = a
        b -= res
        a = 1
    else:
        res = n - b + 1
        a += res
        b = n

    res += (b - a) // 2
    return res

if __name__ == '__main__':
    ans = main(a, b)
    print(ans)
Buggy python submission (verdict: Time Limit Exceeded):
import sys
from bisect import bisect_right as bi_r
from heapq import heapify, heappop, heappush

n, m, v, p, *a = map(int, sys.stdin.read().split())

def main():
    a.sort()
    # return a

    for i in range(n):
        l = bi_r(a, a[i]) - 1
        r = max(0, v - l - 1)
        if r >= p:
            cand = a[l+1:n-p+1]
            ma = cand[-1]
            heapify(cand)
            remain = m * (r - (p - 1))
            while True:
                x = heappop(cand)
                if x == ma:
                    heappush(cand, x)
                    break
                d = ma - x
                if remain < d:
                    remain = 0
                    break
                remain -= d
                heappush(cand, ma)

            if not remain:
                if a[i] + m >= ma:
                    return n - i
            le = len(cand)
            ma += (remain + le - 1) // le
            if a[i] + m >= ma:
                return n - i
        else:
            if a[i] + m >= a[n-p]:
                return n - i

if __name__ == '__main__':
    ans = main()
    print(ans)

Fixed submission:
import sys

n, m, v, p, *a = map(int, sys.stdin.read().split())
a.sort(reverse=True)

def possible(i):
    score = a[i] + m
    if i < p:
        return True
    else:
        every_ok = p - 1 + n - i 
        if every_ok >= v:
            if score >= a[p-1]:
                return True
            else:
                return False
        else:
            s = m * (v - every_ok)
            limit = 0
            for j in range(p-1, i):
                if a[j] > score:
                    return False
                limit += score - a[j]
            if s > limit:
                return False
            return True

def main():
    lo = -1
    hi = n
    while lo + 1 < hi:
        i = (lo + hi) // 2
        if possible(i):
            lo = i
        else:
            hi = i
        
    return lo + 1

if __name__ == '__main__':
    ans = main()
    print(ans)
Buggy python submission (verdict: Time Limit Exceeded):
import sys
import numpy as np 
from scipy.sparse.csgraph import dijkstra
from scipy.sparse import csr_matrix

n, u, v = map(int, sys.stdin.readline().split())
a, b = np.array(sys.stdin.read().split(), dtype=np.int64).reshape(-1, 2).T - 1
graph = csr_matrix(([1] * (n-1), (a, b)), (n, n))

def main():
    dist_u, dist_v = dijkstra(graph, directed=False, indices=[u-1, v-1])
    
    ans = np.amax(dist_v[dist_v >= dist_u]) - 1
    return int(ans)

if __name__ == '__main__':
    ans = main()
    print(ans)

Fixed submission:
import sys

n, u, v, *ab = map(int, sys.stdin.read().split())
ab = list(zip(*[iter(ab)] * 2))
graph = [[] for _ in range(n+1)]
for a, b in ab:
    graph[a].append(b)
    graph[b].append(a)


def main():
    dist_u = [None] * (n + 1)
    dist_u[u] = 0
    stack = [u]

    while stack:
        x = stack.pop()
        for y in graph[x]:
            if dist_u[y] is None:
                dist_u[y] = dist_u[x] + 1
                stack.append(y)

    dist_v = [None] * (n + 1)
    dist_v[v] = 0
    stack = [v]
    contain_v = [False] * (n + 1)
    contain_v[v] = True
    while stack:
        x = stack.pop()
        for y in graph[x]:
            if dist_v[y] is None:
                dist_v[y] = dist_v[x] + 1
                stack.append(y)

    d = 1
    for i in range(1, n+1):
        if dist_v[i] >= dist_u[i]:
            d = max(d, dist_v[i])
    ans = d - 1
    return ans

if __name__ == '__main__':
    ans = main()
    print(ans)