Problem p02616 (python)

Fixed submission:
# -*- coding: utf-8 -*-
import sys

def main():
    N,K = list(map(int, sys.stdin.readline().split()))
    A_list = list(map(int, sys.stdin.readline().split()))


    A_list.sort(key=lambda x: -abs(x))
    mod = 10**9 + 7

    L_plus  = R_plus  = -1 # the Index of plus value
    L_minus = R_minus = -1 # the Index of minus value
    # L : between index 0 to (K-1)
    # R : between index K to end


    for i in range(K-1, -1, -1):
        if (L_plus == -1) and (A_list[i] >= 0):
            L_plus = i
        if (L_minus == -1) and (A_list[i] < 0):
            L_minus = i
        if (L_plus != -1) and (L_minus != -1):
            break
    
    for i in range(K,N):
        if (R_plus == -1) and (A_list[i] >= 0):
            R_plus = i
        if (R_minus == -1) and (A_list[i] < 0):
            R_minus = i
        if (R_plus != -1) and (R_minus != -1):
            break


    cnt_minus = 0

    for i in range(K):
        if A_list[i] < 0:
            cnt_minus += 1


    if cnt_minus % 2 == 0:
        target = A_list[:K]
    
    else: # cnt_minus % 2 != 0
        if (R_plus == -1) and \
           ( (L_plus == -1) or (R_minus == -1) ) :
        
            calc1 = calc2 = 0
        
        elif (R_plus != -1) and \
           ( (L_plus == -1) or (R_minus == -1) ) :
        
            calc1 = 1
            calc2 = 0
        
        elif (R_plus == -1) and \
           ( (L_plus != -1) and (R_minus != -1) ) :
        
            calc1 = 0
            calc2 = 1
        
        elif (R_plus != -1) and \
           ( (L_plus != -1) and (R_minus != -1) ) :
            
            calc1 = A_list[L_plus] * A_list[R_plus]
            calc2 = A_list[L_minus] * A_list[R_minus]


        if calc1 == calc2:
            target = A_list[N-K:N]
        
        elif calc1 > calc2:
            A_list[L_minus], A_list[R_plus] = A_list[R_plus], A_list[L_minus]
            target = A_list[:K]

        elif calc1 < calc2:
            A_list[L_plus], A_list[R_minus] = A_list[R_minus], A_list[L_plus]
            target = A_list[:K]


    ans = 1

    for i in range(len(target)):
        ans *= target[i]
        ans %= mod
    
    print(ans)


if __name__ == "__main__":
    main()

Buggy submission (verdict: Runtime Error):
# -*- coding: utf-8 -*-
import sys

def main():
    N,K = list(map(int, sys.stdin.readline().split()))
    A_list = list(map(int, sys.stdin.readline().split()))

    A_plus = []
    A_minus = []
    mod = 10**9 + 7

    for val in A_list:
        if val >= 0:
            A_plus.append(val)
        else:
            A_minus.append(val)


    A_plus.sort(reverse=True)
    A_minus.sort()

    if (not A_plus):
        ans = 1

        for i in range( len(A_minus)-K, len(A_minus) ):
            ans *= A_minus[i]
            ans %= mod
        
        print(ans)
        return


    if (not A_minus):
        ans = 1

        for i in range(K):
            ans *= A_plus[i]
            ans %= mod
        
        print(ans)
        return


    select = []
    i_p = 0  # index of A_plus
    i_m = 0  # index of A_minus

    while (len(select) < K):
        if (K - len(select)) == 1:
            select.append( A_plus[i_p] )
            i_p += 1
            break

        if (i_m + 1) < len(A_minus):
            if (i_p + 1) < len(A_plus):
                if abs( A_plus[i_p] * A_plus[i_p + 1] ) > abs( A_minus[i_m] * A_minus[i_m + 1] ):
                    select.append( A_plus[i_p] )
                    i_p += 1
                else:
                    select.extend( [ A_minus[i_m] , A_minus[i_m + 1] ] )
                    i_m += 2

            else:
                select.extend( [ A_minus[i_m] , A_minus[i_m + 1] ] )
                i_m += 2
            
        else:
            select.append( A_plus[i_p] )
            i_p += 1
    


    ans = 1
 
    for i in range(K):
        ans *= select[i]
        ans %= mod
    
    print(ans)
    return


if __name__ == "__main__":
    main()

Buggy submission (verdict: Runtime Error):
# -*- coding: utf-8 -*-
import sys

def main():
    N,K = list(map(int, sys.stdin.readline().split()))
    A_list = list(map(int, sys.stdin.readline().split()))

    A_plus = []
    A_minus = []
    mod = 10**9 + 7

    for val in A_list:
        if val >= 0:
            A_plus.append(val)
        else:
            A_minus.append(val)


    A_plus.sort(reverse=True)
    A_minus.sort()

    if (not A_plus):
        ans = 1

        for i in range( len(A_minus)-K, len(A_minus) ):
            ans *= A_minus[i]
            ans %= mod
        
        print(ans)
        return


    if (not A_minus):
        ans = 1

        for i in range(K):
            ans *= A_plus[i]
            ans %= mod
        
        print(ans)
        return


    select = []
    i_p = 0  # index of A_plus
    i_m = 0  # index of A_minus

    while (len(select) < K) and (i_p < len(A_plus)) and (i_m < len(A_minus)):
        if (K - len(select)) == 1:
            select.append( A_plus[i_p] )
            i_p += 1
            break

        if (i_m + 1) < len(A_minus):
            if (i_p + 1) < len(A_plus):
                if abs( A_plus[i_p] * A_plus[i_p + 1] ) > abs( A_minus[i_m] * A_minus[i_m + 1] ):
                    select.append( A_plus[i_p] )
                    i_p += 1
                else:
                    select.extend( [ A_minus[i_m] , A_minus[i_m + 1] ] )
                    i_m += 2

            else:
                select.extend( [ A_minus[i_m] , A_minus[i_m + 1] ] )
                i_m += 2
            
        else:
            select.append( A_plus[i_p] )
            i_p += 1
    


    ans = 1
 
    for i in range(K):
        ans *= select[i]
        ans %= mod
    
    print(ans)
    return


if __name__ == "__main__":
    main()

Buggy submission (verdict: Runtime Error):
# -*- coding: utf-8 -*-
import sys

def main():
    N,K = list(map(int, sys.stdin.readline().split()))
    A_list = list(map(int, sys.stdin.readline().split()))


    A_list.sort(key=lambda x: -abs(x))
    mod = 10**9 + 7

    L_plus  = R_plus  = -1 # the Index of plus value
    L_minus = R_minus = -1 # the Index of minus value
    # L : between index 0 to (K-1)
    # R : between index K to end


    for i in range(K-1, -1, -1):
        if (L_plus == -1) and (A_list[i] >= 0):
            L_plus = i
        if (L_minus == -1) and (A_list[i] < 0):
            L_minus = i
        if (L_plus != -1) and (L_minus != -1):
            break
    
    for i in range(K,N):
        if (R_plus == -1) and (A_list[i] >= 0):
            R_plus = i
        if (R_minus == -1) and (A_list[i] < 0):
            R_minus = i
        if (R_plus != -1) and (R_minus != -1):
            break
    

    if L_plus == R_plus == -1:
        ans = 1

        for i in range( N-K , N ):
            ans *= A_list[i]
            ans %= mod

        print(ans)
        return



    cnt_minus = 0

    for i in range(K):
        if A_list[i] < 0:
            cnt_minus += 1


    if cnt_minus % 2 != 0:
        if (L_minus != -1) and (R_plus != -1):
            calc1 = A_list[L_plus] * A_list[R_plus]
        else:
            calc1 = 0
        
        if (L_plus != -1) and (R_minus != -1):
            calc2 = A_list[L_minus] * A_list[R_minus]
        else:
            calc2 = 0


        if calc1 > calc2:
            A_list[L_minus], A_list[R_plus] = A_list[R_plus], A_list[L_minus]
            target = A_list[:K]

        elif calc1 < calc2:
            A_list[L_plus], A_list[R_minus] = A_list[R_minus], A_list[L_plus]
            target = A_list[:K]
    
    else:
        target = A_list[:K]


    ans = 1

    for i in range(len(target)):
        ans *= target[i]
        ans %= mod
    
    print(ans)


if __name__ == "__main__":
    main()