Problem p02616 (python)

Fixed submission:
def main():
    N, K = map(int, input().split())
    A = list(map(int, input().split()))
    plus = []
    minus = []
    for a in A:
        if a >= 0:
            plus.append(a)
        else:
            minus.append(a)
    # 答えが負になる場合をFalse
    flag = False
    if len(plus) > 0:
        if N == K:
            flag = (len(minus) % 2) == 0
        if N > K:
            flag = True
    elif len(plus) == 0:
        flag = (K % 2) == 0
    
    MOD = 10 ** 9 + 7
    ans = 1
    if flag:
        plus.sort(reverse=True)
        minus.sort()
        if K % 2 == 1:
            ans = plus.pop(0)
        B = []
        for i in range(0,len(plus),2):
            if i + 1 < len(plus):
                B.append(plus[i]*plus[i+1])
        for i in range(0,len(minus),2):
            if i + 1 < len(minus):
                B.append(minus[i]*minus[i+1])
        B.sort(reverse=True)
        for j in range(K//2):
            ans *= B[j]
            ans %= MOD
    else:
        A.sort(reverse=True)
        for i in range(K):
            ans *= A[i]
            ans %= MOD

    print(ans)
        


if __name__ == "__main__":
    main()

Buggy submission (verdict: Runtime Error):
def main():
    N, K = map(int, input().split())
    A = list(map(int, input().split()))
    MOD = 10 ** 9 + 7
    ans = 1
    ok = True
    plus = []
    minus = []
    for a in A:
        if a >= 0:
            ok = False
            plus.append(a)
        else:
            minus.append(a)
    
    if ok and N == K and K % 2 == 0:
        ok = False
    
    if ok:
        A.sort(reverse=True)
        for i in range(K):
            ans *= A[i]
            ans %= MOD
    else:
        plus.sort(reverse=True)
        minus.sort()
        B = []
        if K % 2 == 0:
            for i in range(0,len(plus),2):
                if i + 1 < len(plus):
                    B.append(plus[i] * plus[i+1])
            for i in range(0,len(minus),2):
                if i + 1 < len(minus):
                    B.append(minus[i] * minus[i+1])
            B.sort(reverse=True)
            for i in range(K//2):
                ans *= B[i]
                ans %= MOD
        else:
            ans = plus.pop(0)
            for i in range(0,len(plus),2):
                if i + 1 < len(plus):
                    B.append(plus[i] * plus[i+1])
            for i in range(0,len(minus),2):
                if i + 1 < len(minus):
                    B.append(minus[i] * minus[i+1])
            B.sort(reverse=True)
            for i in range(K//2):
                ans *= B[i]
                ans %= MOD

    print(ans)
            



if __name__ == "__main__":
    main()

Buggy submission (verdict: Runtime Error):
def main():
    N, K = map(int, input().split())
    A = list(map(int, input().split()))
    MOD = 10 ** 9 + 7
    B = []
    minus = [False] * N
    flag = False
    for i in range(N):
        if A[i] == 0:
            continue
        if A[i] < 0:
            minus[i] = True
        if A[i] > 0:
            flag = True
        B.append([abs(A[i]),i])

    if not flag:
        if N - len(B) != 0:
            print(0)
            exit()
        B.sort()
        ans = 1
        for i in range(K):
            ans *= A[B[i][1]]
            ans %= MOD
            
        print(ans)
        exit()

    if K >= len(B):
        if len(B) == N:
            ans = 1
            for i in range(N):
                ans *= A[i]
                ans %= MOD
            print(ans)
            exit()
        else:
            print(0)
            exit()

    B.sort(reverse=True)
    cnt = 0
    cand = []
    min_minus = -1
    min_plus = -1
    for i in range(K):
        if minus[B[i][1]]:
            cnt += 1
            min_minus = B[i][1]
        else:
            min_plus = B[i][1]
        cand.append(B[i][1])
    ans = 1

    if cnt % 2 == 0:
        for i in cand:
            ans *= A[i]
            ans %= MOD
    else:
        if cnt == minus.count(True):
            cand.pop(min_minus)
            for i in range(K,len(B)):
                if B[i][0] > 0:
                    cand.append(B[i][1])
                    break
            else:
                cand.append(min_minus)
            for i in cand:
                ans *= A[i]
                ans %= MOD
        else:
            if A[min_plus] > A[min_minus]:
                cand.pop(min_minus)
                for i in range(K,len(B)):
                    if B[i][0] > 0:
                        cand.append(B[i][1])
                        break
                else:
                    cand.append(min_minus)
            elif A[min_plus] < A[min_minus]:
                cand.pop(min_plus)
                for i in range(K,len(B)):
                    if B[i][0] < 0:
                        cand.append(B[i][1])
                        break
            else:
                cand.pop(max(min_plus,min_minus))
                cand.append(max(min_plus,min_minus) + 1)

            for i in cand:
                ans *= A[i]
                ans %= MOD

            
    print(ans)

    


if __name__ == "__main__":
    main()