Problem p02616 (python)

Fixed submission:
import sys
input = sys.stdin.readline
from itertools import accumulate
from functools import reduce
from operator import mul
from bisect import bisect_left


def read():
    N, K = map(int, input().strip().split())
    A = list(map(int, input().strip().split()))
    return N, K, A


def sign(x):
    if x == 0:
        return 0
    return 1 if x > 0 else -1


def solve(N, K, A, p=10**9+7):
    A.sort(key=lambda a: -abs(a))
    s = 1  # sign
    v = 1  # value
    idx_minus = -1
    idx_not_minus = -1
    for i in range(K):
        _s = sign(A[i])
        s *= _s
        v *= A[i]
        v %= p
        if _s >= 0:
            idx_not_minus = i
        else:
            idx_minus = i
    if s >= 0 or N == K:
        return v
    
    # 以下、左からK個を掛けた数が負になるとき
    # 絶対値が小さな負の数を1つ取り除いて、絶対値が大きな正の数を1つ追加する
    a1 = A[idx_minus]
    b1 = max(A[K:])
    # 絶対値が小さな正の数を1つ取り除いて、絶対値が大きな負の数を1つ追加する
    a2 = A[idx_not_minus]
    b2 = min(A[K:])

    use_1, use_2 = True, True
    if not(idx_minus >= 0 and b1 >= 0):
        use_1 = False
    if not(idx_not_minus >= 0 and b2 < 0):
        use_2 = False
    
    if use_1 and use_2:
        if b1 * a2 >= b2 * a1:
            use_1 = True
            use_2 = False
        else:
            use_1 = False
            use_2 = True
    if use_1:
        v *= pow(a1, p-2, p)
        v %= p
        v *= b1
        v %= p
    elif use_2:
        v *= pow(a2, p-2, p)
        v %= p
        v *= b2
        v %= p
    else:
        # どうやっても正の数が作れないとき
        A.sort(key=lambda a: abs(a))
        s = 1
        v = 1
        for i in range(K):
            _s = sign(A[i])
            s *= _s
            v *= A[i]
            v %= p
    return v


if __name__ == '__main__':
    inputs = read()
    outputs = solve(*inputs)
    if outputs is not None:
        print("%s" % str(outputs))

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline
from itertools import accumulate
from functools import reduce
from operator import mul
from bisect import bisect_left


def read():
    N, K = map(int, input().strip().split())
    A = list(map(int, input().strip().split()))
    return N, K, A


def sign(x):
    if x == 0:
        return 0
    if x < 0:
        return -1
    return 1


class ModInt(int):
    
    MOD = 10**9+7

    def __new__(cls, x, *args, **kwargs):
        return super().__new__(cls, x % ModInt.MOD)

    def __add__(self, other):
        return ModInt(super().__add__(other) % ModInt.MOD)

    def __radd__(self, other):
        return ModInt(super().__radd__(other) % ModInt.MOD)

    def __mul__(self, other):
        return ModInt(super().__mul__(other) % ModInt.MOD)

    def __rmul__(self, other):
        return ModInt(super().__rmul__(other) % ModInt.MOD)

    def __pow__(self, other, *args):
        return ModInt(super().__pow__(other, ModInt.MOD))

    def __rpow__(self, other, *args):
        raise NotImplementedError()

    def __truediv__(self, other):
        inv = ModInt(other).inverse()
        return ModInt(super().__mul__(inv) % ModInt.MOD)

    def __floordiv__(self, other):
        return self.__truediv__(other)

    def __rtruediv__(self, other):
        return ModInt(self.inverse().__mul__(other))

    def __rfloordiv__(self, other):
        return self.__rtruediv__(other)

    def _extgcd(self, a, b):
        assert not isinstance(a, ModInt) and not isinstance(b, ModInt)
        if a < b:
            m, y, x = self._extgcd(b, a)
            return (m, x, y)
        if b == 0:
            x = 1
            y = 0
            return (a, x, y)
        m, y, x = self._extgcd(b, a % b)
        y = y - a // b * x
        return (m, x, y)

    def inverse(self):
        a = int(self)
        m = int(ModInt.MOD)
        n, x, y = self._extgcd(a, m)
        if n != 1:
            raise ValueError("{} and {} are not mutually prime.".format(a, m))
        return ModInt(x)


def solve(N, K, A, p=10**9+7):
    A.sort(key=lambda x: -abs(x))
    base = 1
    sns = 1
    last_minus = -1
    last_plus = -1
    for i in range(K):
        a = A[i]
        base *= a
        base %= p
        sn = sign(a)
        sns *= sn
        if sn < 0:
            last_minus = i
        else:
            last_plus = i
    if sns >= 0:
        return base % p
    
    if last_plus == -1:
        base = 1
        for i in range(K):
            a = A[N-i-1]
            base *= a
            base %= p
        return base % p
    
    base1, base2 = base, base
    # 最も小さい負の数を1つ取り除いて、正の数を1つ足す(絶対値が大きいほど良い)
    assert last_minus < 0 or A[last_minus] < 0
    base1 *= pow(A[last_minus], p-2, p)
    base1 %= p
    m1 = max(A[K:])
    fac1 = abs(1.0 * m1 / A[last_minus])
    has_answer1 = (last_minus > 0 and m1 >= 0)
    base1 *= m1
    base1 %= p
    # 最も小さい"正の数"を1つ取り除いて、負の数を1つ足す(絶対値が大きいほど良い)
    assert last_plus < 0 or A[last_plus] > 0
    base2 *= pow(A[last_plus], p-2, p)
    base2 %= p
    m2 = min(A[K:])
    fac2 = abs(1.0 * m2 / A[last_plus])
    has_answer2 = (last_plus >= 0 and m2 < 0)
    base2 *= m2
    base2 %= p
    if not has_answer1:
        return base2 % p
    elif not has_answer2:
        return base1 % p
    elif fac1 >= fac2:
        return base1 % p
    else:
        return base2 % p



if __name__ == '__main__':
    inputs = read()
    outputs = solve(*inputs)
    if outputs is not None:
        print("%s" % str(outputs))

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline
from itertools import accumulate
from functools import reduce
from operator import mul
from bisect import bisect_left


def read():
    N, K = map(int, input().strip().split())
    A = list(map(int, input().strip().split()))
    return N, K, A


def sign(x):
    if x == 0:
        return 0
    if x < 0:
        return -1
    return 1


class ModInt(int):
    
    MOD = 10**9+7

    def __new__(cls, x, *args, **kwargs):
        return super().__new__(cls, x % ModInt.MOD)

    def __add__(self, other):
        return ModInt(super().__add__(other) % ModInt.MOD)

    def __radd__(self, other):
        return ModInt(super().__radd__(other) % ModInt.MOD)

    def __mul__(self, other):
        return ModInt(super().__mul__(other) % ModInt.MOD)

    def __rmul__(self, other):
        return ModInt(super().__rmul__(other) % ModInt.MOD)

    def __pow__(self, other, *args):
        return ModInt(super().__pow__(other, ModInt.MOD))

    def __rpow__(self, other, *args):
        raise NotImplementedError()

    def __truediv__(self, other):
        inv = ModInt(other).inverse()
        return ModInt(super().__mul__(inv) % ModInt.MOD)

    def __floordiv__(self, other):
        return self.__truediv__(other)

    def __rtruediv__(self, other):
        return ModInt(self.inverse().__mul__(other))

    def __rfloordiv__(self, other):
        return self.__rtruediv__(other)

    def _extgcd(self, a, b):
        assert not isinstance(a, ModInt) and not isinstance(b, ModInt)
        if a < b:
            m, y, x = self._extgcd(b, a)
            return (m, x, y)
        if b == 0:
            x = 1
            y = 0
            return (a, x, y)
        m, y, x = self._extgcd(b, a % b)
        y = y - a // b * x
        return (m, x, y)

    def inverse(self):
        a = int(self)
        m = int(ModInt.MOD)
        n, x, y = self._extgcd(a, m)
        if n != 1:
            raise ValueError("{} and {} are not mutually prime.".format(a, m))
        return ModInt(x)


def solve(N, K, A, p=10**9+7):
    A.sort(key=lambda x: abs(x))
    n_zero = 0
    n_minus = 0
    for a in A:
        if a == 0:
            n_zero += 1
        elif a < 0:
            n_minus += 1
    n_plus = N - n_zero - n_minus

    # 全て0以下のとき
    if n_plus == 0:
        # マイナスを偶数個選べる -> それが最大値
        if n_minus >= K and K % 2 == 0:
            base = 1
            for i in range(K):
                base *= A[N-i]
                base %= p
            return base
        elif n_zero > 0:
            return 0
        else:
            base = 1
            for i in range(K):
                base *= A[i]
                base %= p
            return base % p
    # プラスが1つ以上あるとき
    # 0を選ばざるを得ないなら0
    if n_plus + n_minus < K:
        return 0
    # 0を選ばなくて良いとき: A[n_zero:] から選ぶ
    B = A[n_zero:][::-1]
    sn = 1
    base = 1
    r0 = 0
    r1 = 0
    for i in range(K):
        a = B[i]
        if a < 0:
            r0 = i
            r1 = r0
        sn *= sign(a)
        base *= abs(a)
        base %= p
    if sn == 1:
        return base % p
    
    sn = 1
    base = 1
    for i in range(r1+1):
        a = B[i]
        sn *= sign(a)
        base *= abs(a)
        base %= p
    n = K - r1
    i = r1 + 1
    while n > 0:
        if B[i] > 0:
            base *= a
            base %= p
        n -= 1
    return base % p


if __name__ == '__main__':
    inputs = read()
    outputs = solve(*inputs)
    if outputs is not None:
        print("%s" % str(outputs))

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline
from itertools import accumulate
from functools import reduce
from operator import mul
from bisect import bisect_left


def read():
    N, K = map(int, input().strip().split())
    A = list(map(int, input().strip().split()))
    return N, K, A


def sign(x):
    if x == 0:
        return 0
    if x < 0:
        return -1
    return 1


class ModInt(int):
    
    MOD = 10**9+7

    def __new__(cls, x, *args, **kwargs):
        return super().__new__(cls, x % ModInt.MOD)

    def __add__(self, other):
        return ModInt(super().__add__(other) % ModInt.MOD)

    def __radd__(self, other):
        return ModInt(super().__radd__(other) % ModInt.MOD)

    def __mul__(self, other):
        return ModInt(super().__mul__(other) % ModInt.MOD)

    def __rmul__(self, other):
        return ModInt(super().__rmul__(other) % ModInt.MOD)

    def __pow__(self, other, *args):
        return ModInt(super().__pow__(other, ModInt.MOD))

    def __rpow__(self, other, *args):
        raise NotImplementedError()

    def __truediv__(self, other):
        inv = ModInt(other).inverse()
        return ModInt(super().__mul__(inv) % ModInt.MOD)

    def __floordiv__(self, other):
        return self.__truediv__(other)

    def __rtruediv__(self, other):
        return ModInt(self.inverse().__mul__(other))

    def __rfloordiv__(self, other):
        return self.__rtruediv__(other)

    def _extgcd(self, a, b):
        assert not isinstance(a, ModInt) and not isinstance(b, ModInt)
        if a < b:
            m, y, x = self._extgcd(b, a)
            return (m, x, y)
        if b == 0:
            x = 1
            y = 0
            return (a, x, y)
        m, y, x = self._extgcd(b, a % b)
        y = y - a // b * x
        return (m, x, y)

    def inverse(self):
        a = int(self)
        m = int(ModInt.MOD)
        n, x, y = self._extgcd(a, m)
        if n != 1:
            raise ValueError("{} and {} are not mutually prime.".format(a, m))
        return ModInt(x)


def solve(N, K, A, p=10**9+7):
    A.sort(key=lambda x: -abs(x))
    n_zero = 0
    n_minus = 0
    for a in A:
        if a == 0:
            n_zero += 1
        elif a < 0:
            n_minus += 1
    n_plus = N - n_zero - n_minus

    # 全て0以下のとき
    if n_plus == 0:
        # マイナスを偶数個選べる -> それが最大値
        if n_minus <= K and n_minus % 2 == 0:
            base = 1
            for i in range(K):
                base *= A[N-i]
            return base
        elif n_zero > 0:
            return 0
        else:
            base = 1
            for i in range(K):
                base *= A[i]
            return base % p
    # プラスが1つ以上あるとき
    # 0を選ばざるを得ないなら0
    if n_plus + n_minus < K:
        return 0
    # 0を選ばなくて良いとき: A[n_zero:] から選ぶ
    B = A[n_zero:]
    sn = 1
    base = 1
    r0 = 0
    r1 = 0
    for i in range(K):
        a = A[i]
        if a < 0:
            r0 = i
            r1 = r0
        sn *= sign(a)
        base *= abs(a)
        base %= p
    if sn == 1:
        return base % p
    
    sn = 1
    base = 1
    for i in range(r1):
        a = A[i]
        sn *= sign(a)
        base *= abs(a)
        base %= p
    for i in range(r1, N):
        if A[i] > 0:
            base *= a
            base %= p
    return base % p
        




if __name__ == '__main__':
    inputs = read()
    outputs = solve(*inputs)
    if outputs is not None:
        print("%s" % str(outputs))