Problem p02691 (python)

Fixed submission:
from collections import defaultdict

n = int(input())
a = list(map(int, input().split()))
ap = defaultdict(int)
for i in range(n):
    ap[i - a[i]] += 1

ans = 0
for i in range(n):
    ans += ap[a[i] + i]
print(ans)

Buggy submission (verdict: Runtime Error):
from collections import defaultdict

n, a, b, c = map(int, input().split())
s = [None] * n
for i in range(n):
    s[i] = input().rstrip()

if a + b + c == 0:
    print("No")
    exit(0)

ans = [None] * n
if a + b + c == 1:
    for i in range(n):
        if s[i] == 'AB':
            if a == 0 and b == 0:
                print('No')
                exit(0)
            if b == 0:
                b += 1
                a -= 1
                ans[i] = 'B'
            else:
                a += 1
                b -= 1
                ans[i] = 'A'
        elif s[i] == 'BC':
            if b == 0 and c == 0:
                print('No')
                exit(0)
            if c == 0:
                b -= 1
                c += 1
                ans[i] = 'C'
            else:
                b += 1
                c -= 1
                ans[i] = 'B'
        else:
            if a == 0 and c == 0:
                print('No')
                exit(0)
            if c == 0:
                a -= 1
                c += 1
                ans[i] = 'C'
            else:
                a += 1
                c -= 1
                ans[i] = 'A'

    print("Yes")
    for i in range(n):
        print(ans[i])
    exit(0)

if s[0] == 'AB' and a == 0 and b == 0:
    print("No")
    exit(0)
if s[0] == 'BC' and b == 0 and c == 0:
    print("No")
    exit(0)
if s[0] == 'AC' and a == 0 and c == 0:
    print("No")
    exit(0)

print("Yes")
for i in range(n):
    if s[i] == 'AB':
        if a == 0:
            print('A')
            a += 1
            b -= 1
            continue
        elif b == 0:
            print('B')
            a -= 1
            b += 1
            continue
    elif s[i] == 'BC':
        if b == 0:
            b += 1
            c -= 1
            print('B')
            continue
        elif c == 0:
            b -= 1
            c += 1
            print('C')
            continue
    else:
        if a == 0:
            a += 1
            c -= 1
            print('A')
            continue
        elif c == 0:
            a -= 1
            c += 1
            print('C')
            continue
    
    if a + b + c == 2 and i < n - 1 and s[i] != s[i + 1]:
        if s[i] == 'AB' and a == 1 and b == 1:
            if s[i + 1] == 'AC':
                a += 1
                b -= 1
                print('A')
            else:
                a -= 1
                b += 1
                print('B')
        elif s[i] == 'BC' and b == 1 and c == 1:
            if s[i + 1] == 'AB':
                a -= 1
                b += 1
                print('B')
            else:
                a += 1
                b -= 1
                print('A')
        elif s[i] == 'AC' and a == 1 and c == 1:
            if s[i + 1] == 'AB':
                a += 1
                b -= 1
                print('A')
            else:
                a -= 1
                b += 1
                print('B')
    
    if s[i] == 'AB':
        if a < b:
            a += 1
            b -= 1
            print('A')
        else:
            a -= 1
            b += 1
            print('B')
    elif s[i] == 'BC':
        if b < c:
            b += 1
            c -= 1
            print('B')
        else:
            b -= 1
            c += 1
            print('C')
    else:
        if a < c:
            a += 1
            c -= 1
            print('A')
        else:
            a -= 1
            c += 1
            print('C')