Harryxun/testdata huggingface dataset:
Buggy python submission (verdict: Time Limit Exceeded):
import sys
sys.setrecursionlimit(10**6)


def i1str():
    return input()


def istrs(n=None):
    if n is None:
        return input().split()
    a = []
    for _ in range(n):
        a.append(istrs())
    return a


def i1num():
    return int(input())


def inums(n=None):
    if n is None:
        return list(map(int, input().split()))
    a = []
    for _ in range(n):
        a.append(inums())
    return a


def ostrs(l, sp=" "):
    print(sp.join(l))


def onums(l, sp=" "):
    print(sp.join(map(str, l)))


def solve(aii, index, left, right):
    res = 0
    ai, i = aii[index]
    if index == 0:
        return ai * abs(i-left)

    index -= 1
    # to leftmost
    res1 = ai * abs(i-left) + solve(aii, index, left+1, right)
    # to rightmost
    res2 = ai * abs(right-i) + solve(aii, index, left, right-1)
    return max(res1, res2)


if __name__ == "__main__":
    n = i1num()
    a = inums()
    aii = sorted((a[i], i)for i in range(n))
    print(solve(aii, n-1, 0, n-1))


Fixed submission:
def i1str():
    return input()


def istrs(n=None):
    if n is None:
        return input().split()
    a = []
    for _ in range(n):
        a.append(istrs())
    return a


def i1num():
    return int(input())


def inums(n=None):
    if n is None:
        return list(map(int, input().split()))
    a = []
    for _ in range(n):
        a.append(inums())
    return a


def ostrs(l, sp=" "):
    print(sp.join(l))


def onums(l, sp=" "):
    print(sp.join(map(str, l)))


def solve(n, aii):
    dp = [[0 for _ in range(n+1)] for _ in range(n+1)]

    for x in range(1, n+1):
        ai, i = aii[x-1]
        dp[x][0] = dp[x-1][0] + ai * (i-x+1)
        dp[0][x] = dp[0][x-1] + ai * (n-x-i)

    for s in range(1, n+1):
        for x in range(1, s):
            y = s - x
            ai, i = aii[s-1]
            dp[x][y] = max(dp[x-1][y] + ai * (i-x+1),
                           dp[x][y-1] + ai * (n-y-i))

    res = 0
    for x in range(n+1):
        if dp[x][n-x] > res:
            res = dp[x][n-x]
    return res


if __name__ == "__main__":
    n = i1num()
    a = inums()
    aii = sorted([(a[i], i)for i in range(n)], reverse=True)
    print(solve(n, aii))
Buggy python submission (verdict: Time Limit Exceeded):
from collections import Counter


class Memoize:
    def __init__(self, fn):
        self.fn = fn
        self.memo = {}

    def __call__(self, c, x):
        if x not in self.memo:
            self.memo[x] = self.fn(c, x)
        return self.memo[x]


@Memoize
def cal(c, x):
    s = 0
    for i in c:
        if i == x:
            s += combinations_count(c[i]-1)
        else:
            s += combinations_count(c[i])
    return s


def combinations_count(n):
    return n*(n-1)//2


if __name__ == "__main__":
    n = int(input())
    a = list(map(int, input().split()))
    c = Counter(a)
    for i in range(n):
        print(cal(c, a[i]))


Fixed submission:
from collections import Counter


def combinations_count(n):
    return n*(n-1)//2


if __name__ == "__main__":
    n = int(input())
    a = list(map(int, input().split()))
    c = Counter(a)
    t = 0
    for i in c:
        t += combinations_count(c[i])

    for i in range(n):
        print(t-c[a[i]]+1)
Buggy python submission (verdict: Runtime Error):
a,b = map(int, input().split())
print(int((a-1)/(b-1)+0.5))

Fixed submission:
import math
a,b = map(int, input().split())
print(math.ceil((b-1)/(a-1)))
Buggy python submission (verdict: Runtime Error):
a, b = input().split()
print(max(a+b,a-b,a*b))


Fixed submission:
a, b = map(int,input().split())
print(max(a+b,a-b,a*b))
Buggy python submission (verdict: Runtime Error):
k,x = map(input().split())
start=max(-1000000,x-k+1)
end=min(1000000,x+k-1)
print(" ".join(range(start,end+1)))


Fixed submission:
k, x = map(int, input().split())
start = max(-1000000, x-k+1)
end = min(1000000, x+k-1)
print(" ".join(map(str, range(start, end+1))))
Buggy python submission (verdict: Time Limit Exceeded):
from heapq import heapify, heappush, heappop

n, m = map(int, input().split())

day_value_dict = {}
for _ in range(n):
    a, b = map(int, input().split())
    day_value_dict[a] = day_value_dict.get(a, [])
    heappush(day_value_dict[a], -b)


res = 0
for m_i in range(1, m + 1):
    temp = 0
    selected_x = 0
    for x in range(1, m_i + 1):
        if x not in day_value_dict or not day_value_dict[x]:
            continue
        if day_value_dict[x][0] < temp:
            temp = day_value_dict[x][0]
            selected_x = x
    if selected_x:
        res += temp
        heappop(day_value_dict[selected_x])

print(-res)


Fixed submission:
from heapq import heappush, heappop
n, m = map(int, input().split())

in_dict = {}
for _ in range(n):
    a, b = map(int, input().split())
    in_dict[a] = in_dict.get(a, [])
    in_dict[a].append(b)

heap = []
res = 0
for m_i in range(1, m + 1):
    for i in in_dict.get(m_i, []):
        heappush(heap, -i)
    if heap:
        res += -heappop(heap)

print(res)

Buggy python submission (verdict: Time Limit Exceeded):
n, m = map(int, input().split())
arr = []
for _ in range(n):
    a, b = map(int, input().split())
    arr.append((a, b, 0))


arr.sort(key=lambda x: x[0])
arr.sort(key=lambda x: x[1], reverse=True)

res = 0
for m_i in range(1, m + 1):
    for x in range(len(arr)):
        a, b, deleted = arr[x]
        if deleted:
            continue
        if a <= m_i:
            res += b
            arr[x] = a, b, 1
            break
print(res)


Fixed submission:
from heapq import heappush, heappop
n, m = map(int, input().split())

in_dict = {}
for _ in range(n):
    a, b = map(int, input().split())
    in_dict[a] = in_dict.get(a, [])
    in_dict[a].append(b)

heap = []
res = 0
for m_i in range(1, m + 1):
    for i in in_dict.get(m_i, []):
        heappush(heap, -i)
    if heap:
        res += -heappop(heap)

print(res)

Buggy python submission (verdict: Time Limit Exceeded):
n, m = map(int, input().split())
arr = []
for _ in range(n):
    a, b = map(int, input().split())
    arr.append((a, b))


arr.sort(key=lambda x: x[0])
arr.sort(key=lambda x: x[1], reverse=True)

res = 0
for m_i in range(1, m + 1):
    for x in range(len(arr)):
        a, b = arr[x]
        if a <= m_i:
            res += b
            del arr[x]
            break
print(res)


Fixed submission:
from heapq import heappush, heappop
n, m = map(int, input().split())

in_dict = {}
for _ in range(n):
    a, b = map(int, input().split())
    in_dict[a] = in_dict.get(a, [])
    in_dict[a].append(b)

heap = []
res = 0
for m_i in range(1, m + 1):
    for i in in_dict.get(m_i, []):
        heappush(heap, -i)
    if heap:
        res += -heappop(heap)

print(res)
Buggy python submission (verdict: Runtime Error):
print(sum(range(1,input()+1)))

Fixed submission:
print(sum(range(int(input())+1)))
Buggy python submission (verdict: Runtime Error):
a = int(input())
a = list(map(int,input().split()))
s = 0
for i in range(n):
    s+=1/a[i]
print(1/s)

Fixed submission:
n = int(input())
a = list(map(int,input().split()))
s = 0
for i in range(n):
    s+=1/a[i]
print(1/s)
Buggy python submission (verdict: Runtime Error):
I0 = lambda :map(int,input().split())
I1 = lambda :int(input())
I2 = lambda :list(map(int,input().split()))
#####################################################
A,B = I0()
print(max([A*B,A-B,A+b]))


Fixed submission:
I0 = lambda :map(int,input().split())
I1 = lambda :int(input())
I2 = lambda :list(map(int,input().split()))
#####################################################
A,B = I0()
print(max([A*B,A-B,A+B]))
Buggy python submission (verdict: Runtime Error):
I0 = lambda :map(int,input().split())
I1 = lambda :int(input())
I2 = lambda :list(map(int,input().split()))
#####################################################
from collections import defaultdict
n = I1()
d = defaultdict(int)
for i in range(n):
    d[sorted(input())]+=1
res=0
for i in d.keys():
    if d[i] > 1:
        yy=d[i]-1
        res+=yy*(yy+1)//2
print(res)


Fixed submission:
I0 = lambda :map(int,input().split())
I1 = lambda :int(input())
I2 = lambda :list(map(int,input().split()))
#####################################################
from collections import defaultdict
n = I1()
d = defaultdict(int)
for i in range(n):
    d["".join(sorted(input()))]+=1
res=0
for i in d.keys():
    if d[i] > 1:
        yy=d[i]-1
        res+=yy*(yy+1)//2
print(res)
Buggy python submission (verdict: Runtime Error):
s = [_ for _ in input()]
qn = int(input())
rev = False
prevs = ['']*20000
nexts = ['']*20000
previ = 0
nexti = 0

for _ in range(qn):
  q = [_ for _ in input().split()]
  if q[0] == '1':
    rev = not rev
  else:
    if (rev and q[1] == '2') or (not rev and q[1] == '1'):
      prevs[previ] = q[2]
      previ += 1
    else:
      nexts[nexti] = q[2]
      nexti += 1

prevs.reverse()
s = prevs + s + nexts
if rev:
  s.reverse()
print("".join(s))

Fixed submission:
s = [_ for _ in input()]
qn = int(input())
rev = False
prevs = ['']*200000
nexts = ['']*200000
previ = 0
nexti = 0

for _ in range(qn):
  q = [_ for _ in input().split()]
  if q[0] == '1':
    rev = not rev
  else:
    if (rev and q[1] == '2') or (not rev and q[1] == '1'):
      prevs[previ] = q[2]
      previ += 1
    else:
      nexts[nexti] = q[2]
      nexti += 1

prevs.reverse()
s = prevs + s + nexts
if rev:
  s.reverse()
print("".join(s))

Buggy python submission (verdict: Runtime Error):
s = [_ for _ in input()]
qn = int(input())
rev = False
prevs = ['']*20000
nexts = ['']*20000
previ = 1
nexti = 0

for _ in range(qn):
  q = [_ for _ in input().split()]
  if q[0] == '1':
    rev = not rev
  else:
    if (rev and q[1] == '2') or (not rev and q[1] == '1'):
      prevs[20000-previ] = q[2]
      previ += 1
    else:
      nexts[nexti] = q[2]
      nexti += 1

if rev:
  print("".join(nexts + s + prevs))
else:
  print("".join(prevs + s + nexts))

Fixed submission:
s = [_ for _ in input()]
qn = int(input())
rev = False
prevs = ['']*200000
nexts = ['']*200000
previ = 0
nexti = 0

for _ in range(qn):
  q = [_ for _ in input().split()]
  if q[0] == '1':
    rev = not rev
  else:
    if (rev and q[1] == '2') or (not rev and q[1] == '1'):
      prevs[previ] = q[2]
      previ += 1
    else:
      nexts[nexti] = q[2]
      nexti += 1

prevs.reverse()
s = prevs + s + nexts
if rev:
  s.reverse()
print("".join(s))

Buggy python submission (verdict: Time Limit Exceeded):
s = [_ for _ in input()]
qn = int(input())
rev = False

for _ in range(qn):
  q = [_ for _ in input().split()]
  if q[0] == '1':
    rev = not rev
  else:
    if (rev and q[1] == '2') or (not rev and q[1] == '1'):
      s.insert(0, q[2])
    else:
      s.append(q[2])

if rev:
  s.reverse()
print("".join(s))

Fixed submission:
s = [_ for _ in input()]
qn = int(input())
rev = False
prevs = ['']*200000
nexts = ['']*200000
previ = 0
nexti = 0

for _ in range(qn):
  q = [_ for _ in input().split()]
  if q[0] == '1':
    rev = not rev
  else:
    if (rev and q[1] == '2') or (not rev and q[1] == '1'):
      prevs[previ] = q[2]
      previ += 1
    else:
      nexts[nexti] = q[2]
      nexti += 1

prevs.reverse()
s = prevs + s + nexts
if rev:
  s.reverse()
print("".join(s))
Buggy python submission (verdict: Runtime Error):
import sys
n = input()
arr = [int(_) for _ in input().split()]

for i in arr:
  if i % 2 == 0:
    if i % 3 != 0 and i % 5 != 0:
      print("DENIED")
      sys.exit(1)

print("APPROVED")

Fixed submission:
n = input()
arr = [int(_) for _ in input().split() if int(_) % 2 == 0 and int(_) % 3 != 0 and int(_) % 5 != 0]
 
if len(arr) == 0:
  print("APPROVED")
else:
  print("DENIED")
Buggy python submission (verdict: Time Limit Exceeded):
import sys
import collections

n = input()
arr = [_.replace('\n', '') for _ in sys.stdin]
dic = collections.Counter(arr)

s = [k for k,v in dic.items() if v == max(dic.values())]
for i in sorted(s):
  print(i)


Fixed submission:
import sys
import collections

n = input()
arr = [_.replace('\n', '') for _ in sys.stdin]
dic = collections.Counter(arr)
cnt = list(dic.values()).count(max(dic.values()))
s = [k for k,v in dic.most_common(cnt)]

for i in sorted(s):
  print(i)

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

n = input()
arr = [_.replace('\n', '') for _ in sys.stdin]
dic = {}
for s in arr:
  if s in dic.keys():
    dic[s] += 1
  else:
    dic[s] = 1
s = [k for k,v in dic.items() if v == max(dic.values())]
for i in sorted(s):
  print(i)

Fixed submission:
import sys
import collections

n = input()
arr = [_.replace('\n', '') for _ in sys.stdin]
dic = collections.Counter(arr)
cnt = list(dic.values()).count(max(dic.values()))
s = [k for k,v in dic.most_common(cnt)]

for i in sorted(s):
  print(i)
Buggy python submission (verdict: Time Limit Exceeded):
import numpy as np
n, k = map(int, input().split())
arr = np.array([int(i) for i in input().split()])

max = t = sum(arr[0:k])
max_p = sum([_*(_+1)/(_*2.0) for _ in arr[0:k]])

for i in range(1, n-k+1):
  t += arr[i+k-1]
  t -= arr[i-1]
  if max < t:
    max = t
    max_p = sum([_*(_+1)/(_*2.0) for _ in arr[i:i+k]])

print(max_p)

Fixed submission:
import numpy as np
n, k = map(int, input().split())
arr = np.array([int(_)*(int(_)+1)/(int(_)*2.0) for _ in input().split()])

max = t = sum(arr[0:k])

for i in range(1, n-k+1):
  t += arr[i+k-1]
  t -= arr[i-1]
  if max < t:
    max = t

print(max)

Buggy python submission (verdict: Time Limit Exceeded):
import numpy as np
n, k = map(int, input().split())
arr = np.array([int(i) for i in input().split()])
max = 0
max_p = 0

for i in range(0, n-k+1):
  t = sum(arr[i:i+k])
  if max < t:
    max = t
    max_p = sum([_*(_+1)/(_*2.0) for _ in arr[i:i+k]])

print(max_p)

Fixed submission:
import numpy as np
n, k = map(int, input().split())
arr = np.array([int(_)*(int(_)+1)/(int(_)*2.0) for _ in input().split()])

max = t = sum(arr[0:k])

for i in range(1, n-k+1):
  t += arr[i+k-1]
  t -= arr[i-1]
  if max < t:
    max = t

print(max)

Buggy python submission (verdict: Time Limit Exceeded):
n, k = map(int, input().split())
arr = [int(_) for _ in input().split()]
max = 0
max_p = 0

for i in range(0, n-k+1):
  t = sum(arr[i:i+k])
  if max < t:
    max = t
    max_p = sum([_*(_+1)/(_*2.0) for _ in arr[i:i+k]])

print(max_p)

Fixed submission:
import numpy as np
n, k = map(int, input().split())
arr = np.array([int(_)*(int(_)+1)/(int(_)*2.0) for _ in input().split()])

max = t = sum(arr[0:k])

for i in range(1, n-k+1):
  t += arr[i+k-1]
  t -= arr[i-1]
  if max < t:
    max = t

print(max)

Buggy python submission (verdict: Time Limit Exceeded):
n, k = map(int, input().split())
arr = [int(_)*(int(_)+1)/(int(_)*2.0) for _ in input().split()]
max = 0

for i in range(0, n-k+1):
  t = sum(arr[i:i+k])
  if max < t:
    max = t

print(max)

Fixed submission:
import numpy as np
n, k = map(int, input().split())
arr = np.array([int(_)*(int(_)+1)/(int(_)*2.0) for _ in input().split()])

max = t = sum(arr[0:k])

for i in range(1, n-k+1):
  t += arr[i+k-1]
  t -= arr[i-1]
  if max < t:
    max = t

print(max)
Buggy python submission (verdict: Runtime Error):
a,b,c,d = map(int, input().split(" "))
while True:
  c -= b
  if c <= 0:
    print("Yes")
    break
  a -= d
  if a <= 0:
    print("No")
  	break

Fixed submission:
a,b,c,d = map(int, input().split(" "))
while True:
  c -= b
  if c <= 0:
    print("Yes")
    break
  a -= d
  if a <= 0:
    print("No")
    break
Buggy python submission (verdict: Runtime Error):
_, _, M = map(int, input().split(" "))
a = list(map(int, input().split(" ")))
b = list(map(int, input().split(" ")))
min = min(a) + min(b)
for i in range(M):
  x, y, c = map(int, input().split(" "))
  s = a[x] + b[y] - c
  if min > s:
    min = s
print(min)

Fixed submission:
_, _, M = map(int, input().split(" "))
a = list(map(int, input().split(" ")))
b = list(map(int, input().split(" ")))
min = min(a) + min(b)
for i in range(M):
  x, y, c = map(int, input().split(" "))
  s = a[x-1] + b[y-1] - c
  if min > s:
    min = s
print(min)
Buggy python submission (verdict: Time Limit Exceeded):
n, k = map(int, input().split(" "))
ps = list(map(int, input().split(" ")))

max = 0
for i in range(n-k+1):
  if i == 0:
    s = sum(ps[i:i+k])
  else:
    s = s - (ps[i-1] - ps[i+k-1])
  s = sum(ps[i:i+k])
  if s > max:
    max = s
print((max + k) / 2)

Fixed submission:
n, k = map(int, input().split(" "))
ps = list(map(int, input().split(" ")))

max = 0
for i in range(n-k+1):
  if i == 0:
    s = sum(ps[i:i+k])
  else:
    s = s - (ps[i-1] - ps[i+k-1])
  if s > max:
    max = s
print((max + k) / 2)

Buggy python submission (verdict: Time Limit Exceeded):
n, k = map(int, input().split(" "))
ps = list(map(int, input().split(" ")))

max = 0
for i in range(n-k+1):
  s = sum(ps[i:i+k])
  if s > max:
    max = s
print((max + k) / 2)

Fixed submission:
n, k = map(int, input().split(" "))
ps = list(map(int, input().split(" ")))

max = 0
for i in range(n-k+1):
  if i == 0:
    s = sum(ps[i:i+k])
  else:
    s = s - (ps[i-1] - ps[i+k-1])
  if s > max:
    max = s
print((max + k) / 2)
Buggy python submission (verdict: Runtime Error):
_ = input()
S = input()
print(len(S.split("ABC") - 1))

Fixed submission:
_ = input()
S = input()
print(len(S.split("ABC")) - 1)

Buggy python submission (verdict: Runtime Error):
_ = input()
S = input()
print(len(S.split("ABC") - 1)

Fixed submission:
_ = input()
S = input()
print(len(S.split("ABC")) - 1)
Buggy python submission (verdict: Runtime Error):
S = input()

c = 0
for a, b in zip(s, s[-1]):
  if a == b:
    c += 1

print(len(S)-c)

Fixed submission:
S = input()
l = len(S)
c = 0
for a, b in zip(S[:l//2+1], S[:l//2-1:-1]):
  if a != b:
    c += 1
print(c)
Buggy python submission (verdict: Runtime Error):
n = int(input())
if n / 2 % == 1:
  print("No")
  exit()
s = input()
if s[:n/2] == s[n/2:]:
  print("Yes")
else:
  print("No")

Fixed submission:
n = int(input())
s = input()
if s[:n//2] == s[n//2:]:
  print("Yes")
else:
  print("No")

Buggy python submission (verdict: Runtime Error):
n = int(input())  
s = input()
if s[:n/2] == s[n/2:]:
  print("Yes")
else:
  print("No")

Fixed submission:
n = int(input())
s = input()
if s[:n//2] == s[n//2:]:
  print("Yes")
else:
  print("No")

Buggy python submission (verdict: Runtime Error):
n = int(input())
if n / 2 % == 1:
  print("No")
  return
s = input()
if s[:n/2] == s[n/2:]:
  print("Yes")
else:
  print("No")

Fixed submission:
n = int(input())
s = input()
if s[:n//2] == s[n//2:]:
  print("Yes")
else:
  print("No")
Buggy python submission (verdict: Runtime Error):
from math import sqrt

n = int(input())
for i in range(sqrt(n), 1, -1):
  if n % i == 0:
    print(i + n/i -2)
    break
else:
  print(n-1)

Fixed submission:
from math import sqrt

n = int(input())
for i in range(int(sqrt(n)), 1, -1):
  if n % i == 0:
    print(i + n//i -2)
    break
else:
  print(n-1)
Buggy python submission (verdict: Time Limit Exceeded):
#!/usr/bin/env/ python3
#S = input()
N = int(input())
#A, B, C = input().split()
#N, M = map(int, input().split())

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

#H, W = map(int, input().split())
S = []
count = 0
for i in range(N):
    tmp = input()
    if tmp not in S:
        count += 1
        S.append(tmp)
print(count)


Fixed submission:
#S = input()
N = int(input())
#A, B, C = input().split()
#N, M = map(int, input().split())

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

#H, W = map(int, input().split())
S = [input() for _ in range(N)]
print(len(set(S)))

Buggy python submission (verdict: Time Limit Exceeded):
#!/usr/bin/env/ python3
#S = input()
N = int(input())
#A, B, C = input().split()
#N, M = map(int, input().split())

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

#H, W = map(int, input().split())
S = []
count = 0
for i in range(N):
    tmp = input()
    if tmp not in S:
        count += 1
        S.append(tmp)
print(count)

Fixed submission:
#S = input()
N = int(input())
#A, B, C = input().split()
#N, M = map(int, input().split())

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

#H, W = map(int, input().split())
S = [input() for _ in range(N)]
print(len(set(S)))
Buggy python submission (verdict: Time Limit Exceeded):
n, x, y = map(int, input().split())
lengths = [min(end - start, abs(x-start) + abs(y-end) + 1)
    for start in range(1, n+1)
        for end in range(start+1, n+1)]
for k in range(1, n):
    print(lengths.count(k))

Fixed submission:
n, x, y = map(int, input().split())
lengths = [0 for i in range(n)]
for start in range(1, n+1):
    for end in range(start+1, n+1):
        lengths[min(end - start, abs(x-start) + abs(y-end) + 1)] += 1
print("\n".join(map(str, lengths[1:])))

Buggy python submission (verdict: Time Limit Exceeded):
n, x, y = map(int, input().split())
for k in range(1, n):
    count = 0
    for start in range(1, n+1):
        for end in range(start+1, n+1):
            if min(end - start, abs(x-start) + abs(y-end) + 1) == k:
                count += 1
    print(count)

Fixed submission:
n, x, y = map(int, input().split())
lengths = [0 for i in range(n)]
for start in range(1, n+1):
    for end in range(start+1, n+1):
        lengths[min(end - start, abs(x-start) + abs(y-end) + 1)] += 1
print("\n".join(map(str, lengths[1:])))
Buggy python submission (verdict: Runtime Error):
X, Y, R, G, W = map(int, input().split())
red = sorted(list(map(int, input().split())), key=lambda x: -x)
green = sorted(list(map(int, input().split())), key=lambda x: -x)
white = sorted(list(map(int, input().split())), key=lambda x: -x)

dp = [[[0 for _ in range(X+Y+1)] for __ in range(Y+1)] for ___ in range(X+1)]

ans = 0

for w in range(min(W, X+Y)+1):
    for x in range(X+1):
        for y in range(Y+1):
            if x + y + w == X + Y:
                ans = max(ans, dp[x][y][w])
                break
            if x < X:
                dp[x+1][y][w] = max(dp[x+1][y][w], dp[x][y][w] + red[x])
            if y < Y:
                dp[x][y+1][w] = max(dp[x][y+1][w], dp[x][y][w] + green[y])
            if w < W:
                dp[x][y][w+1] = max(dp[x][y][w+1], dp[x][y][w] + white[w])

print(ans)


Fixed submission:
X, Y, R, G, W = map(int, input().split())
red = list(map(lambda x: (int(x), 0), input().split()))
green = list(map(lambda x: (int(x), 1), input().split()))
white = list(map(lambda x: (int(x), 2), input().split()))

apples = sorted(red + green + white, key=lambda x: x[0])

choice = X + Y
res = 0

while apples and choice:
    val, color = apples.pop()
    if color == 0:
        if X > 0:
            res += val
            X -= 1
            choice -= 1
        continue
    if color == 1:
        if Y > 0:
            res += val
            Y -= 1
            choice -= 1
        continue
    if color == 2:
        res += val
        choice -= 1

print(res)
Buggy python submission (verdict: Time Limit Exceeded):
from collections import Counter
import math

def main():
    r=2
    def comb2(n):
        if n < 2:
            return 0
        return math.factorial(n) // (math.factorial(n - r) * math.factorial(r))
    
    N = int(input())
    A = list(map(int, input().split()))
    
    c = Counter(A)
    
    for ball in A:
        res = 0
        for num in c:
            res += comb2(c[num] - int(num == ball))
        print(res)

if __name__ == "__main__":
    main()

Fixed submission:
from collections import Counter
import math

def main():
    r=2
    def comb2(n):
        if n < 2:
            return 0
        return math.factorial(n) // (math.factorial(n - r) * math.factorial(r))
    
    N = int(input())
    A = list(map(int, input().split()))
    
    count = Counter(A)
    
    all_c = sum([comb2(count[i]) for i in count])
    
    for b in A:
        print(all_c - (count[b]-1))
    
if __name__ == "__main__":
    main()