Problem p02616 (python)

Fixed submission:
N, K = map(int, input().split())
A = list(map(int, input().split()))
MOD = 10 ** 9 + 7

neg_a = sorted([a for a in A if a < 0], reverse=True)
pos_a = sorted([a for a in A if a >= 0])
ans = 1

if len(pos_a) == 0 and K % 2 == 1:
    for a in neg_a[:K]:
        ans = ans * a % MOD
    print(ans)
    exit()

while K > 0:
    if K == 1 or len(neg_a) <= 1:
        if len(pos_a) == 0:
            ans *= neg_a.pop()
        elif len(pos_a) > 0:
            ans *= pos_a.pop()
        K -= 1
    elif len(pos_a) <= 1:
        ans *= neg_a.pop() * neg_a.pop()
        K -= 2
    elif pos_a[-1] * pos_a[-2] > neg_a[-1] * neg_a[-2]:
        ans *= pos_a.pop()
        K -= 1
    else:
        ans *= neg_a.pop() * neg_a.pop()
        K -= 2
    ans %= MOD

print(ans)

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

MOD = 10**9+7
INF = 10**18+10

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

if N == K:
    ans = 1
    for a in A:
        ans *= a
        ans %= MOD
    print(ans)
    exit()

# 個別を読んでいく
not_negative = []
negative = []
for a in A:
    if a < 0:
        negative.append(a)
    else:
        not_negative.append(a)


negR = sorted(negative, reverse=True)

# 全部負で奇数個選ぶとき，答えは負なので最小にする
if len(negative) == N and K % 2 == 1:
    ans = 1
    for k in range(K):
        ans *= negR[k]
        ans %= MOD
    print(ans)
    exit()


not_negative.sort(reverse=True)
negative.sort()

# dequeに入れて2つずつ取る
not_negative = deque(not_negative)
negative = deque(negative)
ans = 1
count = 0

while count < K:
    # あと1つしか使えないとき，負を2つ使えない
    flag1 = count == K - 1
    # 残りの負が1つしかない (2つ使えない)
    flag2 = len(negative) < 2
    # |x1 x2| > |y1 y2| (解法2)
    flag3 = len(not_negative) >= 2 and not_negative[0] * not_negative[1] > negative[0] * negative[1]
    if flag1 or flag2 or flag3:
        ans = ans * not_negative.popleft() % MOD
        count += 1
    else:
        ans = ans * negative.popleft() * negative.popleft() % MOD
        count += 2

print(ans)

Buggy submission (verdict: Runtime Error):
MOD = int(1e9 + 7)

N, K = map(int, input().split())
A = list(map(int, input().split()))
AA = sorted(A, key=lambda x: abs(x))
AAr = sorted(A, key=lambda x: abs(x), reverse=True)


ans = 1
if N == K:
    for a in A:
        ans *= a
        ans %= MOD
elif all([a < 0 for a in A]) and K % 2 == 1:
    for k in range(K):
        ans *= AA[k]
    ans %= MOD
else:
    AArk = AAr[:K]
    AArem = AAr[K:]
    # print(AAr)
    # print(AArk)
    for k in range(K):
        ans *= AAr[k]
    if ans > 0:
        ans %= MOD
    else:
        # 負を1つとり除き，非負を1つ加える
        neg1 = float("inf")
        for a in AArk:
            if a < 0 and abs(a) < abs(neg1):
                neg1 = a
        # 代わりに入れる
        pos1 = None
        for a in AArem:
            if a > 0:
                pos1 = a
                break
        # print(neg1, pos1)

        # 正を1つとり除き，負を1つ加える
        pos2 = float("inf")
        for a in AArk:
            if a > 0 and abs(a) < abs(pos2):
                pos2 = a
        if pos2 == float("inf"):
            ans = ans // neg1 * pos1
            ans %= MOD
        else:
            # 非負にする
            neg2 = None
            for a in AArem:
                if a < 0:
                    neg2 = a
                    break

            # (pos1, neg1) と (pos2, neg2) を比較
            ans1 = ans // neg1 * pos1
            ans2 = ans // pos2 * neg2
            ans = max(ans1, ans2)
            ans %= MOD

print(ans)