Problem p02616 (python)

Fixed submission:
from functools import reduce
from collections import deque
#import operator

#解説見た

def solve(N, K, A):
    MOD = int(1e9 + 7)
    A_plus = [x for x in A if x >= 0]
    A_plus.sort(reverse=True)
    A_minus = [x for x in A if x < 0]
    A_minus.sort()
    #print('+', len(A_plus))
    #print('-',len(A_minus))
    #from collections import Counter
    #print(Counter(A_plus))
    #print(Counter(A_minus))

    if K == 1:
        return max(A) % MOD
    if K == N:
        #ans = reduce(operator.mul, A)
        ans = 1
        for t in A:
            ans *= t
            ans %= MOD
        return ans % MOD

    if len(A_plus) == 0:
        #minusだけから選ぶ場合
        if K % 2 == 1:
            #答えは負になるので、絶対値が小さい順に積を取る
            ans = 1
            for t in list(reversed(A_minus))[:K]:
                ans *= t
                ans %= MOD
            return ans
        else:
            #答えは正になるので、絶対値が大きい順に積をとる
            ans = 1
            for t in A_minus[:K]:
                ans *= t
                ans %= MOD
            return ans
    
    k = K
    S = []
    A_plus = deque(A_plus)
    A_minus = deque(A_minus)
    if k % 2 == 1:
        t = A_plus.popleft()
        S.append(t)
        k -= 1
    #kは偶数になった
    assert k % 2 == 0
    while k > 0:
        if len(A_minus) < 2:
            #print('k',k)
            #print('+', list(A_plus))
            #print('-',list(A_minus))
            t1 = A_plus.popleft()
            S.append(t1)
            t2 = A_plus.popleft()
            S.append(t2)
            k -= 2
        elif len(A_plus) < 2:
            t1 = A_minus.popleft()
            S.append(t1)
            t2 = A_minus.popleft()
            S.append(t2)
            k -= 2
        else:
            t_p = A_plus[0] * A_plus[1]
            t_m = A_minus[0] * A_minus[1]
            if t_p >= t_m:
                t1 = A_plus.popleft()
                S.append(t1)
                t2 = A_plus.popleft()
                S.append(t2)
                assert t1 * t2 == t_p
                k -= 2
            else:
                t1 = A_minus.popleft()
                S.append(t1)
                t2 = A_minus.popleft()
                S.append(t2)
                assert t1 * t2 == t_m
                k -= 2
    #print(Counter(S))
    #print(len(S), K)
    ans = 1
    for t in S:
        ans *= t
        ans %= MOD
    return ans % MOD

if __name__ == "__main__":
    #with open('ABC173E/random_mz03.txt', 'r') as f:
    #    input = f.readline
    N, K = tuple(int(c) for c in input().split())
    A = [int(c) for c in input().split()]
    ret = solve(N, K, A)
    print(ret)

Buggy submission (verdict: Runtime Error):
from functools import reduce
from collections import deque
#import operator

#解説見た

def solve(N, K, A):
    MOD = int(1e9 + 7)
    A_plus = [x for x in A if x >= 0]
    A_plus.sort(reverse=True)
    A_minus = [x for x in A if x < 0]
    A_minus.sort()
    #print('+', len(A_plus))
    #print('-',len(A_minus))
    #from collections import Counter
    #print(Counter(A_plus))
    #print(Counter(A_minus))

    if K == 1:
        return max(A) % MOD
    if K == N:
        #ans = reduce(operator.mul, A)
        ans = 1
        for t in A:
            ans *= t
            ans %= MOD
        return ans % MOD

    if len(A_plus) == 0:
        #minusだけから選ぶ場合
        if K % 2 == 1:
            #答えは負になるので、絶対値が小さい順に積を取る
            ans = 1
            for t in list(reversed(A_minus))[:K]:
                ans *= t
                ans %= MOD
            return ans
        else:
            #答えは正になるので、絶対値が大きい順に積をとる
            ans = 1
            for t in A_minus[:K]:
                ans *= t
                ans %= MOD
            return ans
    
    k = K
    S = []
    A_plus = deque(A_plus)
    A_minus = deque(A_minus)
    if k % 2 == 1:
        t = A_plus.popleft()
        S.append(t)
        k -= 1
    #kは偶数になった
    assert k % 2 == 0
    while k > 0:
        if len(A_minus) < 2:
            #print('k',k)
            #print('+', list(A_plus))
            #print('-',list(A_minus))
            t1 = A_plus.popleft()
            S.append(t1)
            t2 = A_plus.popleft()
            S.append(t2)
            k -= 2
        elif len(A_plus) < 2:
            t1 = A_minus.popleft()
            S.append(t1)
            t2 = A_minus.popleft()
            S.append(t2)
            k -= 2
        else:
            t_p = A_plus[0] * A_plus[1]
            t_m = A_minus[0] * A_minus[1]
            if t_p >= t_m:
                t1 = A_plus.popleft()
                S.append(t1)
                t2 = A_plus.popleft()
                S.append(t2)
                assert t1 * t2 == t_p
                k -= 2
            else:
                t1 = A_minus.popleft()
                S.append(t1)
                t2 = A_minus.popleft()
                S.append(t2)
                assert t1 * t2 == t_m
                k -= 2
    #print(Counter(S))
    #print(len(S), K)
    ans = 1
    for t in S:
        ans *= t
        ans %= MOD
    return ans % MOD

if __name__ == "__main__":
    with open('ABC173E/random_mz03.txt', 'r') as f:
        input = f.readline
        N, K = tuple(int(c) for c in input().split())
        A = [int(c) for c in input().split()]
        ret = solve(N, K, A)
        print(ret)

Buggy submission (verdict: Runtime Error):
from functools import reduce
from collections import deque
#import operator

#解説見た

def solve(N, K, A):
    MOD = int(1e9 + 7)
    A_plus = [x for x in A if x >= 0]
    A_plus.sort(reverse=True)
    A_minus = [x for x in A if x < 0]
    A_minus.sort()
    #print('+', len(A_plus))
    #print('-',len(A_minus))
    #from collections import Counter
    #print(Counter(A_plus))
    #print(Counter(A_minus))

    if K == 1:
        return max(A) % MOD
    if K == N:
        #ans = reduce(operator.mul, A)
        ans = 1
        for t in A:
            ans *= t
            ans %= MOD
        return ans % MOD

    if len(A_plus) == 0:
        #minusだけから選ぶ場合
        if K % 2 == 1:
            #答えは負になるので、絶対値が小さい順に積を取る
            ans = 1
            for t in list(reversed(A_minus))[:K]:
                ans *= t
                ans %= MOD
            return ans
        else:
            #答えは正になるので、絶対値が大きい順に積をとる
            ans = 1
            for t in A_minus[:K]:
                ans *= t
                ans %= MOD
            return ans
    
    k = K
    S = []
    A_plus = deque(A_plus)
    A_minus = deque(A_minus)
    if k % 2 == 1:
        t = A_plus.popleft()
        S.append(t)
        k -= 1
    #kは偶数になった
    assert k % 2 == 0
    while k > 0:
        if len(A_minus) < 2:
            print('k',k)
            print('+', list(A_plus))
            print('-',list(A_minus))
            t1 = A_plus.popleft()
            S.append(t1)
            t2 = A_plus.popleft()
            S.append(t2)
            k -= 2
        elif len(A_plus) < 2:
            t1 = A_minus.popleft()
            S.append(t1)
            t2 = A_minus.popleft()
            S.append(t2)
            k -= 2
        else:
            t_p = A_plus[0] * A_plus[1]
            t_m = A_minus[0] * A_minus[1]
            if t_p >= t_m:
                t1 = A_plus.popleft()
                S.append(t1)
                t2 = A_plus.popleft()
                S.append(t2)
                assert t1 * t2 == t_p
                k -= 2
            else:
                t1 = A_minus.popleft()
                S.append(t1)
                t2 = A_minus.popleft()
                S.append(t2)
                assert t1 * t2 == t_m
                k -= 2
    #print(Counter(S))
    #print(len(S), K)
    ans = 1
    for t in S:
        ans *= t
        ans %= MOD
    return ans % MOD

if __name__ == "__main__":
    #with open('ABC173E/hand_04.txt', 'r') as f:
        #input = f.readline
    N, K = tuple(int(c) for c in input().split())
    A = [int(c) for c in input().split()]
    ret = solve(N, K, A)
    print(ret)

Buggy submission (verdict: Runtime Error):
#import numpy as np
from functools import lru_cache, reduce
import operator


def main():
    MOD = int(1e9 + 7)
    N, K = tuple(int(c) for c in input().split())
    A = [int(c) for c in input().split()]
    A_plus = [x for x in A if x > 0]
    A_plus.sort(reverse=True)
    A_minus = [x for x in A if x < 0]
    A_minus.sort()
    cnt_0 = A.count(0)
    #A_abs = [(abs(x), x) for x in A]
    #A_abs.sort()
    #print(A_plus)
    #print(A_minus)
    #print(A_abs)

    def cum_mat(A, N):
        if len(A) == 0:
            return []
        t = 1
        ret = []
        for i in range(N):
            t *= A[i]
            ret.append(t)
        assert len(ret) == N
        return ret


    #minusからちょうど2つ選んで良い
    A_p_mult = cum_mat(A_plus, max(K, len(A_plus)))
    A_m_mult = cum_mat(A_minus, max(K, len(A_minus)))
    #A_m_pair = [A_minus[i] * A_minus[i+1] for i in range(0,len(A_minus), 2)]
    #print(A_m_pair)

    if K == 1:
        return max(A)

    if len(A_plus) == 0:
        #0またはminusだけから選ぶ場合
        if cnt_0 >= 1:
            return 0
        else:
            if K % 2 == 1:
                #答えは負になるので、絶対値が小さい順に積を取る
                ans = 1
                for t in list(reversed(A_minus))[:K]:
                    ans *= t
                    ans %= MOD
                return ans
            else:
                #答えは正になるので、絶対値が大きい順に積をとる
                ans = 1
                for t in A_minus[:K]:
                    ans *= t
                    ans %= MOD
                return ans

    #答えが正にできない場合で、0を含む場合は0
    if cnt_0 >= 1:
        if len(A_plus) + len(A_minus) < K:
            #0を選ばないといけない場合
            return 0
        elif len(A_plus) + len(A_minus) == K and len(A_minus) % 2 == 1:
            #minusから偶数個取れない場合
            return 0


    @lru_cache(maxsize=None)
    def solve(l: int, m: int):
        #print(l, m)
        if m == 0:
            return A_p_mult[l-1]
        elif l == 0:
            if m % 2 == 0:
                return  A_m_mult[m-1]
            else:
                return reduce(operator.mul, list(reversed(A_minus))[:m])
        if m % 2 == 0:
            p_val = A_p_mult[l-1]
            m_val = A_m_mult[m-1]
            #m_val = A_m_pair[m//2-1]
            return p_val * m_val
        else:
            if ans > -10000000000000000000:
                return -10000000000000000000
            else:
                p_val = reduce(operator.mul, list(reversed(A_plus))[:l])
                m_val = reduce(operator.mul, list(reversed(A_minus))[:m])
                return p_val * m_val
        

    ans = -10000000000000000000
    for l in range(K+1):
        if l > len(A_plus):
            ans = max(ans, solve(len(A_plus), K-len(A_plus)))
        elif K - l > len(A_minus):
            ans = max(ans, solve(K-len(A_minus), len(A_minus)))
        else:
            m = K - l
            ans = max(ans, solve(l, m))
    return ans

if __name__ == "__main__":
    MOD = int(1e9 + 7)
    ret = main()
    ret %= MOD
    print(ret)