Problem p02616 (python)

Fixed submission:
#import numpy as np
Q = int(1e9 + 7)

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

A = sorted(A)

if K == N:
    ans = 1
    for a in A:
        ans *= a
        ans %= Q
    print(ans)
    exit()
elif K == 1:
    ans = A[-1] % Q
    print(ans)
    exit()

A_mi, A_0, A_pl = [], [], []
for a in A:
    if a > 0:
        A_pl.append(a)
    elif a < 0:
        A_mi.append(a)
    else:
        A_0.append(a)

#答えが必ず0
if len(A_mi) + len(A_pl) < K:
    print(0)
    exit()

#全て正
if N == len(A_pl):
    ans, c = 1, 0
    for a in reversed(A_pl):
        ans *= a
        ans %= Q
        c += 1
        if c == K:
            break
    print(ans)
    exit()

#全て負or0
if N == len(A_mi) + len(A_0):
    if K % 2 == 0:
        ans, c = 1, 0
        for a in A_mi:
            ans *= a
            ans %= Q
            c += 1
            if c == K:
                break
        print(ans)
        exit()
    else:
        if len(A_0) > 0:
            print(0)
            exit()
        else:
            ans, c = 1, 0
            for a in reversed(A_mi):
                ans *= a
                ans %= Q
                c += 1
                if c == K:
                    break
            print(ans)
            exit()


cost_2 = []
if len(A_mi) > 1:
    for i in range(0, len(A_mi)-1, 2):
        cost_2.append(A_mi[i] * A_mi[i + 1])

c = 0
ans = 1

if K % 2 == 0:
    while True:
        if len(cost_2) > 0:
            c2 = cost_2[0]
        else:
            c2 = 0
        if len(A_pl) > 1:
            pl2 = A_pl[-1] * A_pl[-2]
        else:
            pl2 = 0
        if c2 > pl2:
            ans *= c2
            cost_2.pop(0)
        else:
            ans *= pl2
            A_pl.pop(-1)
            A_pl.pop(-1)
        ans %= Q
        c += 2
        if c == K:
            break
    print(ans)
else:
    while True:
        if (K - c > 1) and (len(cost_2) > 0):
            c2 = cost_2[0]
            if len(A_pl) > 2:
                pl2 = A_pl[-1] * A_pl[-2]
            else:
                pl2 = 0
            if c2 >= pl2:
                ans *= c2
                cost_2.pop(0)
            else:
                ans *= pl2
                A_pl.pop(-1)
                A_pl.pop(-1)
            ans %= Q
            c += 2
        else:
            ans *= A_pl[-1]
            ans %= Q
            c += 1
            A_pl.pop(-1)
        if K == c:
            break

    print(ans)

Buggy submission (verdict: Runtime Error):
Q = int(1e9 + 7)
N, K = map(int, input().split())
A = list(map(int, input().split()))

A = sorted(A)

if K == N:
    ans = 1
    for a in A:
        ans *= a
        ans %= Q
    print(ans)
    exit()
elif K == 1:
    ans = A[-1] % Q
    print(ans)
    exit()

A_mi, A_0, A_pl = [], [], []
for a in A:
    if a > 0:
        A_pl.append(a)
    elif a < 0:
        A_mi.append(a)
    else:
        A_0.append(a)

#答えが必ず0
if len(A_mi) + len(A_pl) < K:
    print(0)
    exit()

#全て正
if N == len(A_pl):
    ans, c = 1, 0
    for a in reversed(A_pl):
        ans *= a
        ans %= Q
        c += 1
        if c == K:
            break
    print(ans)
    exit()

#全て負or0
if N == len(A_mi) + len(A_0):
    if K % 2 == 0:
        ans, c = 1, 0
        for a in A_mi:
            ans *= a
            ans %= Q
            c += 1
            if c == K:
                break
        print(ans)
        exit()
    else:
        if len(A_0) > 0:
            print(0)
            exit()
        else:
            ans, c = 1, 0
            for a in reversed(A_mi):
                ans *= a
                ans %= Q
                c += 1
                if c == K:
                    break
            print(ans)
            exit()


cost_2 = []
if len(A_mi) > 1:
    for i in range(0, len(A_mi), 2):
        cost_2.append(A_mi[i] * A_mi[i + 1])

c = 0
ans = 1

while True:
    if K - c > 1:
        if len(cost_2) > 0:
            c2 = cost_2[0]
        else:
            c2 = 0
        pl2 = A_pl[-1] * A_pl[-2]
        if c2 >= pl2:
            ans *= c2
            cost_2.pop(0)
        else:
            ans *= pl2
            A_pl.pop(-1)
            A_pl.pop(-1)
        ans %= Q
        c += 2
    else:
        ans *= A_pl[-1]
        ans %= Q
        c += 1
        A_pl.pop(-1)
    if K == c:
        break

print(ans)