Problem p02616 (python)

Fixed submission:
# mを法とするaの乗法的逆元
def modinv(a, mod=10**9+7):
    return pow(a, mod-2, mod)

def main():
    N, K = map(int, input().split())
    A = list(map(int, input().split()))
    MOD = 10**9 + 7

    # もしN = K なら全部選択しないといけないので、
    # 全部掛け算して返す
    if N == K:
        ans = 1
        for a in A:
            ans = (ans * a) % MOD
        print(ans)
        return

    Ap = [] # 正の数だけ入っている
    Am = [] # 負の数だけ入っている.値は絶対値で入れる
    has0 = False # 0がAの中にあるか
    for a in A:
        if a == 0:
            has0 = True
        elif a > 0:
            Ap.append(a)
        else:
            Am.append(-a)

    # それぞれを絶対値の大きい順に並べる
    Ap.sort(reverse=True)
    Am.sort(reverse=True)

    ip = 0
    im = 0
    lp = len(Ap)
    lm = len(Am)

    # もし0でない数がK個未満しかない場合、0を必ず選ばないといけないので
    # 答えは0になる
    if lp + lm < K:
        print(0)
        return

    # # 正の数+負の数=K かつ負の数が奇数のとき、0を選んだほうがマシなので、
    # # 答えは0になる
    # if lp + lm == K and lm % 2 == 1:
    #     print(0)
    #     return
    # もしKが奇数のときで、かつ正の数が1個もない時
    if K % 2 == 1 and lp == 0:
        # 0があるなら0にする
        if has0:
            print(0)
            return
        # 0がない場合は、負の数の中から絶対値の小さいものを選ぶ
        else:
            ans = 1
            for a in Am[-K:]:
                ans = (ans * a) % MOD
            print(MOD - ans)
            return

    # 戦略: Aの中から絶対値の大きい順に選んでいく
    ans = 1
    for i in range(K):
        x = None
        if ip == lp: # 正の数をこれ以上選べないのであれば負の数を選ぶ
            x = Am[im]
            im += 1
        elif im == lm: # 負の数をこれ以上選べないのであれば正の数を選ぶ
            x = Ap[ip]
            ip += 1
        # 正の数、負の数どちらも選べるのであれば、絶対値の大きいほうを選ぶ
        elif Am[im] < Ap[ip]:
            x = Ap[ip]
            ip += 1
        else:
            x = Am[im]
            im += 1
        # print(x)
        ans = ans * x % MOD

    # もし絶対値の大きい順で選んで、負の数を奇数個選んでいたとき
    if im % 2 == 1:
        # 選んだ中で一番小さい負の数を外して、まだ選んでない中で最も大きい正の数を入れるか、
        # 選んだ中で一番小さい正の数を外して、まだ選んでない中で最も（絶対値が）大きい負の数を入れるか選ぶ

        # これ以上0でない数がない場合は、0を選択して答えを0にするのが最大値になる
        if ip == lp and im == lm:
            print(0)
            return
        # 正の数をこれ以上選べない場合は負の数から選ぶ
        elif ip == lp:
            ans = ans * modinv(Ap[ip-1]) % MOD
            ans = ans * Am[im] % MOD
        # 負の数をこれ以上選べない場合は正の数から選ぶ
        # 正の数を1つも選んでないときも正の数から選ぶ
        elif im == lm or ip == 0:
            ans = ans * modinv(Am[im-1]) % MOD
            ans = ans * Ap[ip] % MOD
        # 正、負ともまだ選ぶ余裕がある場合
        else:
            # 一旦、選んだ中で一番小さい負の数と、選んだ中で一番小さい正の数を戻す（再度選べるようにする）
            ip -= 1
            ans = ans * modinv(Ap[ip]) % MOD
            im -= 1
            ans = ans * modinv(Am[im]) % MOD
            
            # 正、負のリストから、まだ選んでない中で絶対値が大きいもの2つを選び、その積の大きいほうを選ぶ
            a = Ap[ip] * Ap[ip+1]
            b = Am[im] * Am[im+1]
            # print(a, b)

            ans = ans * max(a, b) % MOD

    print(ans)

if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
# mを法とするaの乗法的逆元
def modinv(a, mod=10**9+7):
    return pow(a, mod-2, mod)

def main():
    N, K = map(int, input().split())
    A = list(map(int, input().split()))
    MOD = 10**9 + 7
    As = [(abs(a), a<0) for a in A]
    As.sort()
    As = As[::-1]


    # 大きい順
    last_plus = None
    last_minus = None
    is_minus = False
    ans = 1
    for a, m in As[:K]:
        ans = (ans * a) % MOD
        is_minus ^= m
        if m:
            last_minus = a
        else:
            last_plus = a
    
    if last_plus == 0:
        print(0)
        return
    # print(As)
    # print(ans, is_minus)
    # print(last_minus, last_plus)
    # 絶対値の大きい順に選択して、答えが負のとき
    if is_minus:
        # 答えが負の場合はK番目以降の最も絶対値の大きい正負の数を探す
        next_plus = None
        next_minus = None
        has0 = False
        for a, m in As[K:]:
            if a == 0:
                has0 = True
            if m and next_minus == None:
                next_minus = a
            elif not m and next_plus == None and a != 0:
                next_plus = a
        # print(next_minus, next_plus)
        # Aが負の数だけの場合
        if next_plus == None and next_plus == None:
            if has0:
                print(0)
                return
            ans = 1
            for a, m in As[-K:]:
                ans = (ans * a) % MOD
            ans = MOD - ans
        elif next_plus == None:
            if last_plus == None:
                ans = 1
                for a, m in As[-K:]:
                    ans = (ans * a) % MOD
                ans = MOD - ans
            else:
                ans = (ans * modinv(last_plus)) % MOD
                ans = (ans * next_minus) % MOD
        elif next_minus == None:
            ans = (ans * modinv(last_minus)) % MOD
            ans = (ans * next_plus) % MOD
        else:
            a = last_plus * next_plus
            b = last_minus * next_minus
            if a > b:
                ans = (ans * modinv(last_minus)) % MOD
                ans = (ans * next_plus) % MOD
            else:
                ans = (ans * modinv(last_plus)) % MOD
                ans = (ans * next_minus) % MOD

    print(ans)

if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
# mを法とするaの乗法的逆元
def modinv(a, mod=10**9+7):
    return pow(a, mod-2, mod)

def main():
    N, K = map(int, input().split())
    A = list(map(int, input().split()))
    MOD = 10**9 + 7
    As = [(abs(a), a<0) for a in A]
    As.sort()
    As = As[::-1]


    # 大きい順
    last_plus = None
    last_minus = None
    is_minus = False
    ans = 1
    for a, m in As[:K]:
        ans = (ans * a) % MOD
        is_minus ^= m
        if m:
            last_minus = a
        else:
            last_plus = a
    
    if last_plus == 0:
        print(0)
        return
    # print(As)
    # print(ans, is_minus)
    # print(last_minus, last_plus)
    # 絶対値の大きい順に選択して、答えが負のとき
    if is_minus:
        # 答えが負の場合はK番目以降の最も絶対値の大きい正負の数を探す
        next_plus = None
        next_minus = None
        for a, m in As[K:]:
            if m and next_minus == None:
                next_minus = a
            elif not m and next_plus == None:
                next_plus = a
        # print(next_minus, next_plus)
        # Aが負の数だけの場合
        if next_plus == None and next_plus == None:
            ans = 1
            for a, m in As[-K:]:
                ans = (ans * a) % MOD
            ans = MOD - ans
        elif next_plus == None:
            if last_plus == None:
                ans = 1
                for a, m in As[-K:]:
                    ans = (ans * a) % MOD
                ans = MOD - ans
            else:
                ans = (ans * modinv(last_plus)) % MOD
                ans = (ans * next_minus) % MOD
        elif next_minus == None:
            ans = (ans * modinv(last_minus)) % MOD
            ans = (ans * next_plus) % MOD
        else:
            a = last_plus * next_plus
            b = last_minus * next_minus
            if a > b:
                ans = (ans * modinv(last_minus)) % MOD
                ans = (ans * next_plus) % MOD
            else:
                ans = (ans * modinv(last_plus)) % MOD
                ans = (ans * next_minus) % MOD

    print(ans)

if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
# mを法とするaの乗法的逆元
def modinv(a, mod=10**9+7):
    return pow(a, mod-2, mod)

def main():
    N, K = map(int, input().split())
    A = list(map(int, input().split()))
    MOD = 10**9 + 7
    As = [(abs(a), a<0) for a in A]
    As.sort()
    As = As[::-1]


    # 大きい順
    last_plus = None
    last_minus = None
    is_minus = False
    ans = 1
    for a, m in As[:K]:
        ans = (ans * a) % MOD
        is_minus ^= m
        if m:
            last_minus = a
        else:
            last_plus = a
    
    # print(As)
    # print(ans, is_minus)
    # print(last_minus, last_plus)
    # 絶対値の大きい順に選択して、答えが負のとき
    if is_minus:
        # 答えが負の場合はK番目以降の最も絶対値の大きい正負の数を探す
        next_plus = None
        next_minus = None
        for a, m in As[K:]:
            if m and next_minus == None:
                next_minus = a
            elif not m and next_plus == None:
                next_plus = a
        # print(next_minus, next_plus)
        # Aが負の数だけの場合
        if next_plus == None and next_plus == None:
            ans = 1
            for a, m in As[-K:]:
                ans = (ans * a) % MOD
            ans = MOD - ans
        elif next_plus == None:
            if last_plus == None:
                ans = 1
                for a, m in As[-K:]:
                    ans = (ans * a) % MOD
                ans = MOD - ans
            else:
                ans = (ans * modinv(last_plus)) % MOD
                ans = (ans * next_minus) % MOD
        elif next_minus == None:
            ans = (ans * modinv(last_minus)) % MOD
            ans = (ans * next_plus) % MOD
        else:
            a = last_plus * next_plus
            b = last_minus * next_minus
            if a > b:
                ans = (ans * modinv(last_minus)) % MOD
                ans = (ans * next_plus) % MOD
            else:
                ans = (ans * modinv(last_plus)) % MOD
                ans = (ans * next_minus) % MOD

    print(ans)

if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
# mを法とするaの乗法的逆元
def modinv(a, mod=10**9+7):
    return pow(a, mod-2, mod)

def main():
    N, K = map(int, input().split())
    A = list(map(int, input().split()))
    MOD = 10**9 + 7
    As = [(abs(a), a<0) for a in A]
    As.sort()
    As = As[::-1]


    # 大きい順
    last_plus = None
    last_minus = None
    is_minus = False
    ans = 1
    for a, m in As[:K]:
        ans = (ans * a) % MOD
        is_minus ^= m
        if m:
            last_minus = a
        else:
            last_plus = a
    
    # print(As)
    # print(ans, is_minus)
    # print(last_minus, last_plus)
    # 絶対値の大きい順に選択して、答えが負のとき
    if is_minus:
        # 答えが負の場合はK番目以降の最も絶対値の大きい正負の数を探す
        next_plus = None
        next_minus = None
        for a, m in As[K:]:
            if m and next_minus == None:
                next_minus = a
            elif not m and next_plus == None:
                next_plus = a
        # print(next_minus, next_plus)
        # Aが負の数だけの場合
        if next_plus == None:
            if last_plus == None:
                ans = 1
                for a, m in As[-K:]:
                    ans = (ans * a) % MOD
                ans = MOD - ans
            else:
                ans = (ans * modinv(last_plus)) % MOD
                ans = (ans * next_minus) % MOD
        elif next_minus == None:
                ans = (ans * modinv(last_minus)) % MOD
                ans = (ans * next_plus) % MOD
        else:
            a = last_plus * next_plus
            b = last_minus * next_minus
            if a > b:
                ans = (ans * modinv(last_minus)) % MOD
                ans = (ans * next_plus) % MOD
            else:
                ans = (ans * modinv(last_plus)) % MOD
                ans = (ans * next_minus) % MOD

    print(ans)

if __name__ == '__main__':
    main()