Problem p02616 (python)

Fixed submission:
mod = 10**9+7
N,K = map(int,input().split())
A = list(map(int,input().split()))
A.sort(key=lambda x:-abs(x))

if N==K:
    ans = 1
    for i in range(N):
        ans *= A[i]
        ans %= mod
elif A[K-1]==0:
    ans = 0
else:
    p = []
    p_ = []
    n = []
    n_ = []
    for i in range(N):
        a = A[i]
        if i<K:
            if a>0:
                p.append(a)
            elif a<0:
                n.append(a)
        else:
            if a>0:
                p_.append(a)
            elif a<0:
                n_.append(a)
    if len(n)%2==0:
        ans = 1
        for a in n:
            ans *= a
            ans %= mod
        for a in p:
            ans *= a
            ans %= mod
    else:
        if len(p)+len(p_)==0:
            if len(n)+len(n_)==N:
                n += n_
                ans = 1
                for i in range(K):
                    ans *= n[-i-1]
                    ans %= mod
            else:
                ans = 0
        else:
            if len(p)==0:
                n.pop()
                ans = p_[0]
                for a in n:
                    ans *= a
                    ans %= mod
            else:
                cand = []
                ap = p.pop()
                if p_:
                    cand.append([ap,p_[0]])
                an = n.pop()
                if n_:
                    cand.append([an,n_[0]])
                if len(cand)==2:
                    ap1,ap2 = cand[0]
                    an1,an2 = cand[1]
                    ans = max(ap1*ap2,an1*an2)%mod
                    for a in n:
                        ans *= a
                        ans %= mod
                    for a in p:
                        ans *= a
                        ans %= mod
                else:
                    ans = cand[0][0]*cand[0][1]%mod
                    for a in n:
                        ans *= a
                        ans %= mod
                    for a in p:
                        ans *= a
                        ans %= mod
print(ans)

Buggy submission (verdict: Runtime Error):
mod = 10**9+7
N,K = map(int,input().split())
A = list(map(int,input().split()))
A.sort(key=lambda x:-abs(x))
p = []
n = []
for i in range(K):
    a = A[i]
    if a>0:
        p.append(a)
    elif a<0:
        n.append(a)

if len(n)%2==0:
    ans = 1
    for a in n:
        ans *= a
        ans %= mod
    for a in p:
        ans *= a
        ans %= mod
else:
    ap = 0
    an = 0
    for i in range(K,N):
        if ap==0 and A[i]>0:
            ap = A[i]
        if an==0 and A[i]<0:
            an = -A[i]
    if len(p)>0:
        ap_ = p[-1]
    else:
        ap_ = 1
    an_ = -n[-1]
    if ap==0 and an==0:
        ans = 0
    else:
        if ap*ap_>an*an_:
            n.pop()
            p.append(ap)
        else:
            p.pop()
            n.append(-an)
        ans = 1
        for a in n:
            ans *= a
            ans %= mod
        for a in p:
            ans *= a
            ans %= mod

if len(p)+len(n)<K:
    ans = 0
if len(p)==0 and len(n)==K and K%2==1:
    An = [A[i] for i in range(N) if A[i]<0]
    An.sort()
    An = An[::-1]
    ans = 1
    for i in range(K):
        ans *= An[i]
        ans %= mod
    if A[-1]==0:
        ans = 0
if N==K:
    ans = 1
    for i in range(K):
        ans *= A[i]
        ans %= mod
print(ans%mod)

Buggy submission (verdict: Runtime Error):
mod = 10**9+7
N,K = map(int,input().split())
A = list(map(int,input().split()))
A.sort(key=lambda x:-abs(x))
p = []
n = []
for i in range(K):
    a = A[i]
    if a>0:
        p.append(a)
    else:
        n.append(a)

if len(n)%2==0:
    ans = 1
    for a in n:
        ans *= a
        ans %= mod
    for a in p:
        ans *= a
        ans %= mod
else:
    ap = 0
    an = 0
    for i in range(K,N):
        if ap==0 and A[i]>0:
            ap = A[i]
        if an==0 and A[i]<0:
            an = A[i]
    if len(p)*len(n)>0:
        ap_ = p[-1]
        an_ = -n[-1]
        if ap_*an_==0:
            ans = 0
        else:
            if ap*ap_>an*an_:
                n.pop()
                p.append(ap)
            else:
                p.pop()
                n.append(an)
            ans = 1
            for a in n:
                ans *= a
                ans %= mod
            for a in p:
                ans *= a
                ans %= mod
    elif len(p)>0:
        ans = 1
        for a in p:
            ans *= a
            ans %= mod
    elif len(n)>0:
        for a in n:
            ans *= a
            ans %= mod
if len(p)+len(n)<K:
    ans = 0
if len(p)==0 and len(n)==K and K%2==1:
    An = [A[i] for i in range(N) if A[i]<0]
    An.sort()
    An = An[::-1]
    ans = 1
    for i in range(K):
        ans *= An[i]
        ans %= mod
print(ans)