Harryxun/testdata huggingface dataset:
Buggy python submission (verdict: Runtime Error):
#include <iostream>
#include <algorithm>
#include <vector>


int main()
{
    long long  n,m,k; std::cin >> n >> m >> k;
    // long long  a[n+1], b[m+1];
    std::vector<long long>  a{n+1}, b{m+1};
    a[0]=0; b[0]=0;
    for (int i=1;i<n+1;i++)
        std::cin >> a[i];
    for (int i=1; i<m+1; i++)
        std::cin >> b[i];
    
    for (int i=1; i<n+1; i++)
        a[i]+=a[i-1];
    for (int i=1; i<m+1; i++)
        b[i]+=b[i-1];

    int ans=0, r=m;
    for (int i=0; i<n+1; i++)
    {
        if(a[i]>k) break;
        while (a[i]+b[r]>k)
            r--;
        ans = std::max({ans, i+r});
    }
    
    std::cout << ans;

    return 0;
}

Fixed submission:
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines

n,m,k = map(int, readline().split())
a = list(map(int, readline().split()))
b = list(map(int, readline().split()))

sa = [0]*(n+1)
sb = [0]*(m+1)

for i in range(n):
    sa[i+1] = sa[i]+a[i]
for i in range(m):
    sb[i+1] = sb[i]+b[i]

isb=m
ans = 0
for isa in range(n+1):
    if sa[isa]>k:
        break
    #while isb>0 and sa[isa]+sb[isb]>k:
    while sa[isa]+sb[isb]>k:
        isb-=1
    ans = max(ans, isa+isb)
    
print(ans)

Buggy python submission (verdict: Runtime Error):
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines

n,m,k = map(int, readline().split())
a = list(map(int, readline().split()))
b = list(map(int, readline().split()))

sa = [0]*(n+1)
sb = [0]*(m+1)

for i in range(n):
    sa[i+1] = sa[i]+a[i]
for i in range(m):
    sb[i+1] = sb[i]+b[i]

isb=m
ans = 0
for isa in range(n+1):
    if sa[isa]>k:
        break
    while isb>0 and sa[isa]+sb[isb]>K:
        isb-=1
    ans = max(ans, isa+isb)
    
print(ans)

Fixed submission:
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines

n,m,k = map(int, readline().split())
a = list(map(int, readline().split()))
b = list(map(int, readline().split()))

sa = [0]*(n+1)
sb = [0]*(m+1)

for i in range(n):
    sa[i+1] = sa[i]+a[i]
for i in range(m):
    sb[i+1] = sb[i]+b[i]

isb=m
ans = 0
for isa in range(n+1):
    if sa[isa]>k:
        break
    #while isb>0 and sa[isa]+sb[isb]>k:
    while sa[isa]+sb[isb]>k:
        isb-=1
    ans = max(ans, isa+isb)
    
print(ans)
Buggy python submission (verdict: Runtime Error):
n = input()
s = [input() for i in range(n)]
u = set(s)
dic = dict(zip(u, [0]*len(u)))
for i in range(n):
    dic[s[i]] += 1
m = max(dic.values())
res = []
for k, v in dic.items():
    if v == m:
        res.appned(k)
res.sort()
print(*res, sep="\n")

Fixed submission:
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
N = int(readline())
S = read().split()

S = sorted(S)
ls = [[] for i in range(N + 1)]
max_num = 0
left = 0
right = 0
while right < N:
    while right < N and S[left] == S[right]:
        right += 1
    if (right  - left) >= max_num:
        ls[right - left].append(S[left].decode())
        max_num = right - left
    left = right
    
print('\n'.join(ls[max_num]))

Buggy python submission (verdict: Runtime Error):
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
N = int(readline())
S = read().split()

S = sorted(S)
ls = [[] for i in range(N + 1)]
max_num = 0
left = 0
right = 0
while right < N:
    while right < N and S[left] == S[right]:
        right += 1
    if (right  - left) >= max_num:
        ls[right - left].append(S[left])
        max_num = right - left
    left = right
    
print('\n'.join(ls[max_num]))

Fixed submission:
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
N = int(readline())
S = read().split()

S = sorted(S)
ls = [[] for i in range(N + 1)]
max_num = 0
left = 0
right = 0
while right < N:
    while right < N and S[left] == S[right]:
        right += 1
    if (right  - left) >= max_num:
        ls[right - left].append(S[left].decode())
        max_num = right - left
    left = right
    
print('\n'.join(ls[max_num]))

Buggy python submission (verdict: Time Limit Exceeded):
from operator import itemgetter 

import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
N = int(readline())
S = read().split()

dic = {}
for i in S:
    if i in dic:
        dic[i] += 1
    else:
        dic[i] = 1
        
dic_sorted = sorted(dic.items(), key=itemgetter(1, 0))[::-1]

res = []
n = dic_sorted[0][1]
i = 0
l = len(dic_sorted)
while i < l and dic_sorted[i][1] == n:
    #print(dic_sorted[i][0])
    res = [dic_sorted[i][0].decode()] + res
    i += 1
    
print('\n'.join(res))

Fixed submission:
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
N = int(readline())
S = read().split()

S = sorted(S)
ls = [[] for i in range(N + 1)]
max_num = 0
left = 0
right = 0
while right < N:
    while right < N and S[left] == S[right]:
        right += 1
    if (right  - left) >= max_num:
        ls[right - left].append(S[left].decode())
        max_num = right - left
    left = right
    
print('\n'.join(ls[max_num]))
Buggy python submission (verdict: Runtime Error):
from itertools import permutations as perm
import sys
import math
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
N = int(readline())
XY = np.array(read().split(), np.int64)
X = XY[0::2]
Y = XY[1::2]

res = 0
cnt = 0
for p in perm(range(N)):
    for i in range(N - 1):
        res += math.sqrt((X[p[i + 1]] - X[p[i]])**2 + (Y[p[i + 1]] - Y[p[i]])**2)
    cnt += 1    
print(res/cnt)

Fixed submission:
import numpy as np
import sys
import math
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
N = int(readline())
XY = np.array(read().split(), np.int64)
X = XY[0::2]
Y = XY[1::2]

t = 0
for i in range(N):
    for j in range(i, N):
        t += math.sqrt((X[i] - X[j])**2 + (Y[i] - Y[j])**2)
print(2*t/N)
Buggy python submission (verdict: Runtime Error):
N = int(input())
D = list(map(int, input().split()))
if D[0] != 0:
    print(0)
else:
    D = D[1:]
    C = [0]*N
    for d in D:
        C[d] += 1
    if C[0] != 0:
        print(0)
    else:
        res = 1
        for i in range(1, N - 1):
            res *= C[i]**C[i + 1]
    print(res%998244353)

Fixed submission:
N = int(input())
D = list(map(int, input().split()))
if D[0] != 0:
    print(0)
else:
    D = D[1:]
    C = [0]*N
    m = max(D)
    for d in D:
        C[d] += 1
    if C[0] != 0 or 0 in C[1: m + 1]:
        print(0)
    else:
        res = 1
        i = 1
        for i in range(1, m):
            res *= (C[i]**C[i + 1])%998244353
        print(res%998244353)

Buggy python submission (verdict: Runtime Error):
N = int(input())
D = list(map(int, input().split()))
if D[0] != 0:
    print(0)
else:
    D = D[1:]
    C = [0]*N
    for d in D:
        C[d] += 1
    if C[0] != 0:
        print(0)
    else:
        res = 1
        for i in range(1, N - 1):
            res *= C[i]**C[i + 1]
    print(res)

Fixed submission:
N = int(input())
D = list(map(int, input().split()))
if D[0] != 0:
    print(0)
else:
    D = D[1:]
    C = [0]*N
    m = max(D)
    for d in D:
        C[d] += 1
    if C[0] != 0 or 0 in C[1: m + 1]:
        print(0)
    else:
        res = 1
        i = 1
        for i in range(1, m):
            res *= (C[i]**C[i + 1])%998244353
        print(res%998244353)
Buggy python submission (verdict: Time Limit Exceeded):
import numpy as np
A, B = map(int, input().split())
while A%B > 0:
    c = A
    A = B
    B = c%B
P = [1]
if B%2 == 0:
    P.append(2)
for i in range(3, B//2 + 1, 2):
    if B%i == 0:
        P.append(i)
if B > 2:
    P.append(B)
    
l = len(P)
for i in range(2, l):
    if P[i] == 0:
        continue
    for j in range(i + 1, l):
        if P[j]%P[i] == 0:
            P[j] = 0
P = np.array(P)
print(np.where(P != 0)[0].size)

Fixed submission:
def gcd(x, y):
    """
    ユークリッドの互除法を用いてxとyの最大公約数を求める。
    """
    while True:
        r = x%y
        if r == 0:
            break
        x = y
        y = r
    return y
    
def prime_factorize(x):
    """
    素因数分解
    """
    res = []
    while x%2 == 0:
        res.append(2)
        x = x//2
    p = 3
    while p*p <= x:
        while x%p == 0:
            res.append(p)
            x = x//p
        p += 2
    if x != 1:
        res.append(x)
    return res


A, B = map(int, input().split())
print(len(set(prime_factorize(gcd(A, B)))) + 1)

Buggy python submission (verdict: Runtime Error):
mport numpy as np
A, B = map(int, input().split())
while A%B > 0:
    c = A
    A = B
    B = c%B
P = [1]
if B%2 == 0:
    P.append(2)
for i in range(3, B//2 + 1, 2):
    if B%i == 0:
        P.append(i)
if B > 2:
    P.append(B)
    
l = len(P)
for i in range(2, l):
    if P[i] == 0:
        continue
    for j in range(i + 1, l):
        if P[j]%P[i] == 0:
            P[j] = 0
P = np.array(P)
print(np.where(P != 0)[0].size)

Fixed submission:
def gcd(x, y):
    """
    ユークリッドの互除法を用いてxとyの最大公約数を求める。
    """
    while True:
        r = x%y
        if r == 0:
            break
        x = y
        y = r
    return y
    
def prime_factorize(x):
    """
    素因数分解
    """
    res = []
    while x%2 == 0:
        res.append(2)
        x = x//2
    p = 3
    while p*p <= x:
        while x%p == 0:
            res.append(p)
            x = x//p
        p += 2
    if x != 1:
        res.append(x)
    return res


A, B = map(int, input().split())
print(len(set(prime_factorize(gcd(A, B)))) + 1)
Buggy python submission (verdict: Time Limit Exceeded):
import numpy as np
N, K, Q = map(int, input().split())
point = np.zeros(N, dtype='int') + K
for i in range(Q):
    x = int(input())
    point[x - 1] = point[x - 1] + 1
    point = point - 1
    
for i in range(N):
    if point[i] > 0:
        print('Yes')
    else:
        print('No')

Fixed submission:
N, K, Q = map(int, input().split())
d = K - Q
point = [0]*N
for i in range(Q):
    point[int(input()) - 1] += 1
    
for i in range(N):
    if point[i] + d > 0:
        print('Yes')
    else:
        print('No')
Buggy python submission (verdict: Runtime Error):
def resolve():
    N = int(input())
    A = list(map(int, input().split()))
    x = []
    for i in range(1, N):
        x.append(sum(A[i:]))

    ans = 0
    for i in range(N-1):
        ans += A[i] * x[i]
    ans = ans % (10**9 + 7)
    print(ans)

if __name__ == "__main__":
    resolve()

Fixed submission:
def resolve():
    D, T, S = map(int, input().split())
    if D <= T * S:
        print("Yes")
    else:
        print("No")

if __name__ == "__main__":
    resolve()
Buggy python submission (verdict: Time Limit Exceeded):
def resolve():
    N = int(input())
    A = list(map(int, input().split()))
    x = []
    for i in range(1, N):
        x.append(sum(A[i:]))

    ans = 0
    for i in range(N-1):
        ans += A[i] * x[i]
    ans = ans % (10**9 + 7)
    print(ans)

if __name__ == "__main__":
    resolve()

Fixed submission:
def resolve():
    N = int(input())
    A = list(map(int, input().split()))
    x = [0] * N
    tot = 0
    for i in range(N-1, 0, -1):
        tot += A[i]
        x[i] = tot

    ans = 0
    for i in range(N-1):
        ans += A[i] * x[i+1]
    ans = ans % (10**9 + 7)
    print(ans)

if __name__ == "__main__":
    resolve()

Buggy python submission (verdict: Time Limit Exceeded):
def resolve():
    N = int(input())
    A = list(map(int, input().split()))
    x = []
    for i in range(1, N):
        x.append(sum(A[i:]))

    ans = 0
    for i in range(N-1):
        ans += A[i] * x[i]
    ans = ans % (10**9 + 7)
    print(ans)

if __name__ == "__main__":
    resolve()

Fixed submission:
def resolve():
    N = int(input())
    A = list(map(int, input().split()))
    x = [0] * N
    tot = 0
    for i in range(N-1, 0, -1):
        tot += A[i]
        x[i] = tot

    ans = 0
    for i in range(N-1):
        ans += A[i] * x[i+1]
    ans = ans % (10**9 + 7)
    print(ans)

if __name__ == "__main__":
    resolve()

Buggy python submission (verdict: Time Limit Exceeded):
def resolve():
    N = int(input())
    A = list(map(int, input().split()))
    m = 10**9 + 7

    ans = 0
    for i in range(N-1):
        for j in range(i+1, N):
            ans += A[i] * A[j]
    ans = ans % m
    print(ans)

if __name__ == "__main__":
    resolve()

Fixed submission:
def resolve():
    N = int(input())
    A = list(map(int, input().split()))
    x = [0] * N
    tot = 0
    for i in range(N-1, 0, -1):
        tot += A[i]
        x[i] = tot

    ans = 0
    for i in range(N-1):
        ans += A[i] * x[i+1]
    ans = ans % (10**9 + 7)
    print(ans)

if __name__ == "__main__":
    resolve()
Buggy python submission (verdict: Time Limit Exceeded):
def resolve():
    N, K = map(int, input().split())
    P = list(map(int, input().split()))
    C = list(map(int, input().split()))
    for i in range(N):
        P[i] -= 1

    ans = -1e18
    for i in range(N):
        n = P[i]
        score = C[n]
        t = [score]
        while n != i:
            n = P[n]
            score += C[n]
            t.append(score)

        total = t[-1]
        l = len(t)
        for i in range(min(K, l)):
            temp_score = t[i]
            if total > 0 and K > l:
                e = int((K-(i+1))/l)
                temp_score += total*e
            ans = max(ans, temp_score)

    print(ans)


if __name__ == "__main__":
    resolve()

Fixed submission:
def resolve():
    N, K = map(int, input().split())
    P = list(map(int, input().split()))
    C = list(map(int, input().split()))
    for i in range(N):
        P[i] -= 1

    ans = -1e18
    for i in range(N):
        n = P[i]
        score = C[n]
        t = [score]
        while n != i:
            n = P[n]
            score += C[n]
            t.append(score)

        total = t[-1]
        l = len(t)
        for i in range(min(K, l)):
            temp_score = t[i]
            if total > 0 and K > l:
                e = int((K-(i+1))/l)
                temp_score += total*e
            ans = max(ans, temp_score)

    print(ans)


if __name__ == "__main__":
    resolve()

Buggy python submission (verdict: Runtime Error):
def resolve():
    N, K = map(int, input().split())
    P = list(map(int, input().split()))
    C = list(map(int, input().split()))

    ans = -100100100
    for i in range(N):
        node_set = set()
        n = P[i]
        score = 0
        t = []
        while n not in node_set:
            score += C[n-1]
            t.append(score)
            node_set.add(n)
            n = P[n-1]

        if K < len(t):
            max_t = max(t[:K])
        else:
            if t[-1] > 0:
                temp = t[-1] * int(K / len(t)) + max([0] + t[:K%len(t)])
                max_t = max(max_t, temp)
            else:
                max_t = max(t)
        ans = max(ans, max_t)

    print(ans)


if __name__ == "__main__":
    resolve()

Fixed submission:
def resolve():
    N, K = map(int, input().split())
    P = list(map(int, input().split()))
    C = list(map(int, input().split()))
    for i in range(N):
        P[i] -= 1

    ans = -1e18
    for i in range(N):
        n = P[i]
        score = C[n]
        t = [score]
        while n != i:
            n = P[n]
            score += C[n]
            t.append(score)

        total = t[-1]
        l = len(t)
        for i in range(min(K, l)):
            temp_score = t[i]
            if total > 0 and K > l:
                e = int((K-(i+1))/l)
                temp_score += total*e
            ans = max(ans, temp_score)

    print(ans)


if __name__ == "__main__":
    resolve()
Buggy python submission (verdict: Runtime Error):
from functools import reduce
import operator

NK = list(map(int, input().split(' ')))
N = NK[0]
K = NK[1]
A_list = list(map(int, input().split(' ')))
score_list = []

for g in range(K, N+1):
    score = reduce(operator.mul, A_list[g-K: g])
    score_list.append(score)

for i in range(len(score_list) - 1):
    if (score_list[i] < score_list[i+1]):
        print('Yes')
    else:
        print('No')

Fixed submission:
rate = int(input())
k = '1'
if (400 <= rate < 600):
    k = '8'
elif (600 <= rate < 800):
    k = '7'
elif (800 <= rate < 1000):
    k = '6'
elif (1000 <= rate < 1200):
    k = '5'
elif (1200 <= rate < 1400):
    k = '4'
elif (1400 <= rate < 1600):
    k = '3'
elif (1600 <= rate < 1800):
    k = '2'
elif (1800 <= rate < 2000):
    k = '1'
print(k)
Buggy python submission (verdict: Time Limit Exceeded):
def resolve():
    N = int(input())
    A = list(map(int, input().split()))

    def feelGood(table: list):
        pos = value = 0
        if len(table) == 0:
            return pos, value
        circle = table + [table[0]]
        for i in range(len(circle)-1):
            v = min(circle[i], circle[i+1])
            if (v > value):
                pos = i
                value = v
        return pos, value

    A.sort(reverse=True)
    table = []
    total = 0

    for i in A:
        pos, value = feelGood(table)
        table.insert(pos, i)
        total += value
    print(total)

if __name__ == "__main__":
    resolve()

Fixed submission:
import math

def resolve():
    N = int(input())
    A = list(map(int, input().split()))

    A.sort()

    total = 0
    for i in range(1, N):
        total += A[N - 1 - math.floor(i/2)]
    print(total)

if __name__ == "__main__":
    resolve()
Buggy python submission (verdict: Time Limit Exceeded):
def resolve():
    M, N, K = map(int, input().split())
    a = list(map(int, input().split()))
    b = list(map(int, input().split()))

    count = 0
    for m in range(M+1):
        sumA = sum(a[:m])
        if sumA > K:
            break
        for n in range(N+1):
            sumB = sum(b[:n])
            if sumB > K:
                break
            if (sumA + sumB) <= K:
                count = max(count, m + n)
            else:
                break
    print(count)

if __name__ == "__main__":
    resolve()

Fixed submission:
def resolve():
    N, M, K = map(int, input().split())
    A = list(map(int, input().split()))
    B = list(map(int, input().split()))

    a, b = [0], [0]
    for i in range(N):
        a.append(a[i] + A[i])
    for i in range(M):
        b.append(b[i] + B[i])

    count = 0
    j = M
    for i in range(N + 1):
        if a[i] > K:
            break
        while b[j] > K - a[i]:
            j -= 1
        count = max(count, i + j)
    print(count)

if __name__ == "__main__":
    resolve()
Buggy python submission (verdict: Time Limit Exceeded):
def resolve():
    N = int(input())

    ans = 0
    for n in range(1, N+1):
        for m in range(n, N+1, n):
            ans += m
    print(ans)

if __name__ == "__main__":
    resolve()

Fixed submission:
def resolve():
    N = int(input())

    ans = 0
    for i in range(1, N+1):
        n = N // i
        ans += i * (n + 1) * n // 2
    print(ans)

if __name__ == "__main__":
    resolve()

Buggy python submission (verdict: Time Limit Exceeded):
def resolve():
    N = int(input())

    def calc(x: int, n: int):
        y = int(n/x)
        return int(y * (y + 1) * x / 2)

    result = 0
    for i in range(1, N+1):
        result += calc(i, N)
    print(result)

if __name__ == "__main__":
    resolve()

Fixed submission:
def resolve():
    N = int(input())

    ans = 0
    for i in range(1, N+1):
        n = N // i
        ans += i * (n + 1) * n // 2
    print(ans)

if __name__ == "__main__":
    resolve()
Buggy python submission (verdict: Runtime Error):
def resolve():
    N = int(input())

    alp = [0] * 6
    for i in range(6):
        x = 26 ** (5-i)
        alp[i] = N // x
        N = N % x
    ans = ''
    for a in alp:
        if a > 0:
            ans += chr(64+a)
    print(ans.lower())

if __name__ == "__main__":
    resolve()

Fixed submission:
def resolve():
    N = int(input())

    ans = ''
    while N > 0:
        N -= 1
        ans += chr(ord("a") + N % 26)
        N = N // 26
    print(ans[::-1])

if __name__ == "__main__":
    resolve()
Buggy python submission (verdict: Time Limit Exceeded):
def resolve():
    N = int(input())
    A = list(map(int, input().split()))
    Q = int(input())
    BC = [list(map(int, input().split())) for _ in range(Q)]

    a_dict = {}
    for i in A:
        if i in a_dict:
            a_dict[i] += 1
        else:
            a_dict[i] = 1

    for b, c in BC:
        if b in a_dict:
            if c in a_dict:
                a_dict[c] += a_dict.pop(b)
            else:
                a_dict[c] = a_dict.pop(b)
        print(sum([k * v for k, v in a_dict.items()]))

if __name__ == "__main__":
    resolve()

Fixed submission:
def resolve():
    N = int(input())
    A = list(map(int, input().split()))
    Q = int(input())
    BC = [list(map(int, input().split())) for _ in range(Q)]

    a = [0] * 100001
    for i in A:
        a[i] += i

    ans = sum(a)
    for b, c in BC:
        if a[b] == 0:
            print(ans)
            continue
        move = a[b] // b
        a[b] = 0
        a[c] += c * move
        ans += (c - b) * move
        print(ans)

if __name__ == "__main__":
    resolve()
Buggy python submission (verdict: Runtime Error):
def resolve():
    N = int(input())
    A = list(map(int, input().split()))
    A.sort()
    a = [0] * 100005
    A_max = A[-1]

    for i in A:
        for j in range(1, A_max//i + 1):
            a[i * j] += 1

    ans = 0
    for i in A:
        if a[i] == 1:
            ans += 1
    print(ans)

if __name__ == "__main__":
    resolve()

Fixed submission:
def resolve():
    N = int(input())
    A = list(map(int, input().split()))
    A.sort()
    a = [0] * 1000005
    A_max = A[-1]

    for i in A:
        for j in range(1, A_max//i + 1):
            a[i * j] += 1

    ans = 0
    for i in A:
        if a[i] == 1:
            ans += 1
    print(ans)

if __name__ == "__main__":
    resolve()

Buggy python submission (verdict: Runtime Error):
def resolve():
    N = int(input())
    A = list(map(int, input().split()))
    A.sort()
    a = [0] * 10005
    A_max = A[-1]

    for i in A:
        for j in range(1, A_max//i + 1):
            a[i * j] += 1

    ans = 0
    for i in A:
        if a[i] == 1:
            ans += 1
    print(ans)

if __name__ == "__main__":
    resolve()

Fixed submission:
def resolve():
    N = int(input())
    A = list(map(int, input().split()))
    A.sort()
    a = [0] * 1000005
    A_max = A[-1]

    for i in A:
        for j in range(1, A_max//i + 1):
            a[i * j] += 1

    ans = 0
    for i in A:
        if a[i] == 1:
            ans += 1
    print(ans)

if __name__ == "__main__":
    resolve()
Buggy python submission (verdict: Time Limit Exceeded):
def resolve():
    N = int(input())
    A = list(map(int, input().split()))
    A.insert(0, 0)

    ans = 0
    for i in range(1, N+1):
        for j in range(A[i]+i+1, N+1):
            if j - i == A[j] + A[i]:
                ans += 1
    print(ans)

if __name__ == "__main__":
    resolve()

Fixed submission:
def resolve():
    N = int(input())
    A = list(map(int, input().split()))

    mp = dict()
    ans = 0
    for i in range(N):
        x = i - A[i]
        ans += mp.get(x, 0)
        y = A[i] + i
        mp[y] = mp.get(y, 0) + 1
    print(ans)

if __name__ == "__main__":
    resolve()

Buggy python submission (verdict: Time Limit Exceeded):
def resolve():
    N = int(input())
    A = list(map(int, input().split()))
    A.insert(0, 0)

    ans = 0
    for i in range(1, N+1):
        for j in range(A[i]+i+1, N+1):
            if j - i == A[j] + A[i]:
                ans += 1
    print(ans)

if __name__ == "__main__":
    resolve()

Fixed submission:
def resolve():
    N = int(input())
    A = list(map(int, input().split()))

    mp = dict()
    ans = 0
    for i in range(N):
        x = i - A[i]
        ans += mp.get(x, 0)
        y = A[i] + i
        mp[y] = mp.get(y, 0) + 1
    print(ans)

if __name__ == "__main__":
    resolve()