Problem p02616 (python)

Fixed submission:
def main():
    N, K = map(int, input().split())
    A = list(map(int, input().split()))

    plus = [A[i] for i in range(N) if A[i] >= 0]
    minus = [A[i] for i in range(N) if A[i] < 0]

    ans = 1

    P = len(plus)
    M = len(minus)
    plus.sort(reverse=True)
    minus.sort()
    mod = 10**9 + 7

    if N == K:
        for x in A:
            ans *= x
            ans %= mod
    elif P == 0:
        if K % 2 == 1:
            minus.sort(reverse=True)
        for i in range(K):
            ans *= minus[i]
            ans %= mod
    else:  # 必ず正の値が作れる？ 1つ以上余るので、偶奇は選べるはず
        id = 0
        if K % 2 == 1:
            ans = plus[0]
            id = 1
            P -= 1
            K -= 1
        pp = [plus[i] * plus[i+1] for i in range(id, P-1+id, 2)]
        mm = [minus[i] * minus[i+1] for i in range(0, M-1, 2)]

        id = 0
        id2 = 0

        for i in range(K // 2):
            if id < P // 2 and (id2 >= M // 2 or pp[id] >= mm[id2]):
                ans *= pp[id]
                id += 1
            else:
                ans *= mm[id2]
                id2 += 1
            ans %= mod

    print(ans % mod)


main()

Buggy submission (verdict: Runtime Error):
def main():
    N, K = map(int, input().split())
    A = list(map(int, input().split()))

    plus = [A[i] for i in range(N) if A[i] >= 0]
    minus = [A[i] for i in range(N) if A[i] < 0]

    ans = 1

    id = 0
    id2 = 0
    P = len(plus)
    M = len(minus)
    plus.sort(reverse=True)

    mod = 10**9 + 7

    if P == 0:
        if K % 2 == 0:
            minus.sort()
            for i in range(K):
                ans *= minus[i]
                ans %= mod
        else:
            minus.sort(reverse=True)
            for i in range(K):
                ans *= minus[i]
                ans %= mod

    elif M == 0:
        for i in range(K):
            ans *= plus[i]
            ans %= mod

    else:

        p = P // 2
        m = M // 2
        if K > (p+m)*2 and (K % 2 == 1 and P % 2 == 0):  # 絶対値が小さい方から貪欲
            A.sort(key=lambda x: abs(x))
            for i in range(K):
                ans *= A[i]
                ans %= mod

        else:  # 絶対値が大きい方から貪欲
            minus.sort()
            pp = [plus[i]*plus[i+1] for i in range(M // 2)]
            mm = [minus[i] * minus[i+1] for i in range(P//2)]

            id = 0
            id2 = 0

            pp.sort(reverse=True)
            mm.sort(reverse=True)
            pp.append(1)
            mm.append(1)
            for i in range((K // 2) - 1):
                if id == p - 1:
                    ans *= mm[id2]
                    id2 += 1
                elif id2 == m:
                    ans *= pp[id]
                    id += 1
                elif pp[id] >= mm[id2]:
                    ans *= pp[id]
                    id += 1
                else:
                    ans *= mm[id2]
                    id2 += 1

                ans %= mod

            if K % 2 == 1:

                if id >= p-1:
                    ans *= mm[id2]
                    ans *= plus[id*2]
                else:
                    ans *= pp[id]
                    id += 1
                    ans *= pp[id*2]
            else:
                if id == p - 1:
                    ans *= mm[id2]
                    id2 += 1
                elif id2 == m:
                    ans *= pp[id]
                    id += 1
                elif pp[id] >= mm[id2]:
                    ans *= pp[id]
                    id += 1
                else:
                    ans *= mm[id2]
                    id2 += 1

    print(ans % mod)


main()

Buggy submission (verdict: Runtime Error):
def main():
    N, K = map(int, input().split())
    A = list(map(int, input().split()))

    plus = [A[i] for i in range(N) if A[i] >= 0]
    minus = [A[i] for i in range(N) if A[i] < 0]

    ans = 1

    id = 0
    id2 = 0
    P = len(plus)
    M = len(minus)
    plus.sort(reverse=True)

    mod = 10**9 + 7

    if P == 0:
        if K % 2 == 0:
            minus.sort()
            for i in range(K):
                ans *= minus[i]
                ans %= mod
        else:
            minus.sort(reverse=True)
            for i in range(K):
                ans *= minus[i]
                ans %= mod

    elif M == 0:
        for i in range(K):
            ans *= plus[i]
            ans %= mod

    else:

        p = P // 2
        m = M // 2
        if K > (p+m)*2 and (K % 2 == 1 and P % 2 == 0):  # 絶対値が小さい方から貪欲
            A.sort(key=lambda x: abs(x))
            for i in range(K):
                ans *= A[i]
                ans %= mod

        else:  # 絶対値が大きい方から貪欲
            minus.sort()
            pp = [plus[i]*plus[i+1] for i in range(M // 2)]
            mm = [minus[i] * minus[i+1] for i in range(P//2)]

            id = 0
            id2 = 0

            pp.sort(reverse=True)
            mm.sort(reverse=True)

            for i in range((K // 2) - 1):
                if id == p - 1:
                    ans *= mm[id2]
                    id2 += 1
                elif id2 == m:
                    ans *= pp[id]
                    id += 1
                elif pp[id] >= mm[id2]:
                    ans *= pp[id]
                    id += 1
                else:
                    ans *= mm[id2]
                    id2 += 1

                ans %= mod

            if K % 2 == 1:
                if id == p-1:
                    ans *= mm[id2]
                    ans *= plus[id*2]
                else:
                    ans *= pp[id]
                    id += 1
                    ans *= pp[id*2]
            else:
                if id == p - 1:
                    ans *= mm[id2]
                    id2 += 1
                elif id2 == m:
                    ans *= pp[id]
                    id += 1
                elif pp[id] >= mm[id2]:
                    ans *= pp[id]
                    id += 1
                else:
                    ans *= mm[id2]
                    id2 += 1

    print(ans % mod)


main()

Buggy submission (verdict: Runtime Error):
def main():
    N, K = map(int, input().split())
    A = list(map(int, input().split()))

    plus = [A[i] for i in range(N) if A[i] >= 0]
    minus = [A[i] for i in range(N) if A[i] < 0]

    ans = 1

    id = 0
    id2 = 0
    P = len(plus)
    M = len(minus)
    plus.sort(reverse=True)

    mod = 10**9 + 7

    if P == 0:
        if K % 2 == 0:
            minus.sort()
            for i in range(K):
                ans *= minus[i]
                ans %= mod
        else:
            minus.sort(reverse=True)
            for i in range(K):
                ans *= minus[i]
                ans %= mod

    elif M == 0:
        for i in range(K):
            ans *= plus[i]
            ans %= mod

    else:

        p = P // 2
        m = M // 2
        if K > (p+m)*2 and (K % 2 == 1 and P % 2 == 0):  # 絶対値が小さい方から貪欲
            A.sort(key=lambda x: abs(x))
            for i in range(K):
                ans *= A[i]
                ans %= mod

        else:  # 絶対値が大きい方から貪欲
            minus.sort()
            pp = [plus[i]*plus[i+1] for i in range(M // 2)]
            mm = [minus[i] * minus[i+1] for i in range(P//2)]

            id = 0
            id2 = 0

            pp.sort(reverse=True)
            mm.sort(reverse=True)

            for i in range(K // 2):
                if id == p:
                    ans *= mm[id2]
                    id2 += 1
                elif id2 == m:
                    ans *= pp[id]
                    id += 1

                elif pp[id] >= mm[id2]:
                    ans *= pp[id]
                    id += 1
                else:
                    ans *= mm[id2]
                    id2 += 1

                ans %= mod

            if K % 2 == 1:
                ans *= plus[id * 2]

    print(ans % mod)


main()