Problem p02616 (python)

Fixed submission:
from heapq import heappush, heappop


N, K = map(int, input().split())
A = list(map(int, input().split()))
A = sorted(A, reverse=True)

MOD = 10**9 + 7

# case1. N == K のときは答えは確定
if N == K:
    ans = 1
    for i in range(K):
        ans *= A[i]
        ans %= MOD
    print(ans)
# case2. A が全部負の数のとき、 Kが奇数ならば残念ながら答えは負になるのでなるべく絶対値を小さくする
elif max(A) < 0 and K % 2 == 1:
    ans = 1
    for i in range(K):
        ans *= A[i]
        ans %= MOD
    print(ans)
else:
    # case3. それ以外は正にできる

    # K が奇数のとき、非負要素から1つ取って偶数の場合に帰着
    # K が偶数のとき、面倒なので 2 つセットで考えて K//2 個取る問題に帰着

    if K%2 == 1:
        ans = A.pop(0)  # あらかじめ最大値を確保
        n = (K-1)//2
    else:
        ans = 1
        n = K//2

    # 絶対値の降順にソートする
    A = sorted(A, key=lambda x: abs(x), reverse=True)

    # 候補テーブルを作成
    cand = []  # 非負要素、負要素をそれぞれ絶対値の大きい順に 2 つずつ掛けたもの

    BIGINT = 10**9 + 1  # tmp 用
    tmp1 = BIGINT
    tmp2 = BIGINT

    for a in A:
        if a >= 0:
            if tmp1 == BIGINT:
                tmp1 = a
            else:
                heappush(cand, -tmp1*a)
                tmp1 = BIGINT
        else:
            if tmp2 == BIGINT:
                tmp2 = a
            else:
                heappush(cand, -tmp2*a)
                tmp2 = BIGINT

    # あとは大きい方から n 個取る
    for i in range(n):
        ans *= -heappop(cand)
        ans %= MOD
    print(ans)

Buggy submission (verdict: Runtime Error):
from heapq import heappush, heappop


N, K = map(int, input().split())
A = list(map(int, input().split()))
A = sorted(A, reverse=True)

MOD = 10**9 + 7

# case1. N == K のときは答えは確定
if N == K:
    ans = 1
    for i in range(K):
        ans *= A[i]
        ans %= MOD
    print(ans)
# case2. A が全部負の数のとき、 Kが奇数ならば残念ながら答えは負になるのでなるべく絶対値を小さくする
elif max(A) < 0 and K % 2 == 1:
    ans = 1
    for i in range(K):
        ans *= A[i]
        ans %= MOD
    print(ans)
else:
    # case3. それ以外は正にできる

    # K が奇数のとき、非負要素から1つ取って偶数の場合に帰着
    # K が偶数のとき、面倒なので 2 つセットで考えて K//2 個取る問題に帰着

    if K%2 == 1:
        ans = A.pop(0)  # あらかじめ最大値を確保
        n = (K-1)//2
    else:
        ans = 1
        n = K//2

    # 絶対値の降順にソートする
    A = sorted(A, key=lambda x: abs(x), reverse=True)

    # 候補テーブルを作成
    cand = []  # 非負要素、負要素をそれぞれ絶対値の大きい順に 2 つずつ掛けたもの

    BIGINT = 10**9 + 1  # tmp 用
    tmp1 = BIGINT
    tmp2 = BIGINT
    for a in A:
        if a >= 0:
            if tmp1 == BIGINT:
                tmp1 = A
            else:
                heappush(cand, -tmp1*a)
                tmp1 = BIGINT
        else:
            if tmp2 == BIGINT:
                tmp2 = a
            else:
                heappush(cand, -tmp2*a)
                tmp2 = BIGINT

    # あとは大きい方から n 個取る
    for i in range(n):
        ans *= -heappop(cand)
        ans %= MOD
    print(ans)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
A = list(map(int, input().split()))
MOD = 10**9 + 7

# case1. N == K のときは答えは確定
if N == K:
    ans = 1
    for i in range(K):
        ans *= A[i]
        ans %= MOD
    print(ans)
# case2. A が全部負の数のとき、 Kが奇数ならば残念ながら答えは負になるのでなるべく絶対値を小さくする
elif max(A) < 0 and K % 2 == 1:
    ans = 1
    A = sorted(A, reverse=True)
    for i in range(K):
        ans *= A[i]
        ans %= MOD
    print(ans)
else:
    # case3. それ以外は正にできるので、絶対値の降順にソートする
    A = sorted(A, key=lambda x: abs(x), reverse=True)

    # 候補テーブルを作成
    A1 = []  # 非負要素のリスト（降順）
    A2 = []  # 負要素を絶対値の大きい順に2つずつ掛けたリスト（降順）
    tmp = 1
    for a in A:
        if a >= 0:
            A1.append(a)
        else:
            if tmp != 1:
                A2.append(tmp*a)
                tmp = 1
            else:
                tmp *= a

    # cost が K を上回らなように大きいものから埋めていく〜
    cost = 0
    while cost < K:
        if (K - cost) >= 2:
            # このときは A2 からも入れてよし
            if A1 and A2:
                if A1[0] >= A2[0]:
                    cost += 1
                    ans *= A1.pop(0)
                else:
                    cost += 2
                    ans *= A2.pop(0)
            elif A1:
                cost += 1
                ans *= A1.pop(0)
            else:
                cost += 2
                ans *= A2.pop(0)
            ans %= MOD
        else:
            cost += 1
            ans *= A1.pop(0)
            ans %= MOD

    print(ans)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
A = list(map(int, input().split()))
MOD = 10**9 + 7

# case1. N == K のときは答えは確定
if N == K:
    ans = 1
    for i in range(K):
        ans *= A[i]
        ans %= MOD
    print(ans)
# case2. A が全部負の数のとき、 Kが奇数ならば残念ながら答えは負になるのでなるべく絶対値を小さくする
elif max(A) < 0 and K % 2 == 1:
    ans = 1
    A = sorted(A, reverse=True)
    for i in range(K):
        ans *= A[i]
        ans %= MOD
    print(ans)
else:
    # case3. それ以外は正にできるので、絶対値の降順にソートする
    A = sorted(A, key=lambda x: abs(x), reverse=True)
    # とりあえず大きい順に K 個掛けてみる
    ans = 1
    for i in range(K):
        ans *= A[i]
    if ans >= 0:
        ans %= MOD
        print(ans)
    else:
        # 取り除く対象を old, 取り込む対象を new とする
        # case3-1. 負数を取り除いて非負数を取り込む
        # まずは new を探す
        j = K
        while j < N and A[j] < 0:
            j += 1
        if j < N:
            new1 = A[j]
        # old を探す
        l = K-1
        while A[l] >= 0:
            l -= 1
        old1 = A[l]
        if new1 and old1:
            cand1 = ans // old1 * new1
        # case3-2. 非負数を取り除いて負数を取り込む
        # まずは new を探す
        j = K
        while j < N and A[j] >= 0:
            j += 1
        if j < N:
            new2 = A[j]
        # old を探す
        l = K-1
        while A[l] < 0:
            l -= 1
        old2 = A[l]
        if new2 and old2:
            cand2 = ans // old2 * new2

        if cand1 and cand2:
            ans = max(cand1, cand2) % MOD
        elif cand1:
            ans = cand1 % MOD
        else:
            ans = cand2 % MOD
        print(ans)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
A = list(map(int, input().split()))
MOD = 10**9 + 7

# case1. N == K のときは答えは確定
if N == K:
    ans = 1
    for i in range(K):
        ans *= A[i]
        ans %= MOD
    print(ans)
# case2. A が全部負の数のとき、 Kが奇数ならば残念ながら答えは負になるのでなるべく絶対値を小さくする
elif max(A) < 0 and K % 2 == 1:
    ans = 1
    A = sorted(A, reverse=True)
    for i in range(K):
        ans *= A[i]
        ans %= MOD
    print(ans)
else:
    # case3. それ以外は正にできるので、絶対値の降順にソートする
    A = sorted(A, key=lambda x: abs(x), reverse=True)
    # とりあえず大きい順に K 個掛けてみる
    ans = 1
    for i in range(K):
        ans *= A[i]
    if ans >= 0:
        ans %= MOD
        print(ans)
    else:
        # 取り除く対象を old, 取り込む対象を new とする
        # case3-1. 負数を取り除いて非負数を取り込む
        # まずは new を探す
        j = K
        while j < N or A[j] < 0:
            j += 1
        if j < N:
            new1 = A[j]
        # old を探す
        l = K-1
        while A[l] >= 0:
            l -= 1
        old1 = A[l]
        if new1 and old1:
            cand1 = ans // old1 * new1
        # case3-2. 非負数を取り除いて負数を取り込む
        # まずは new を探す
        j = K
        while j < N or A[j] >= 0:
            j += 1
        if j < N:
            new2 = A[j]
        # old を探す
        l = K-1
        while A[l] < 0:
            l -= 1
        old2 = A[l]
        if new2 and old2:
            cand2 = ans // old2 * new2

        if cand1 and cand2:
            print(cand1, cand2)
            ans = max(cand1, cand2) % MOD
        elif cand1:
            ans = cand1 % MOD
        else:
            ans = cand2 % MOD
        print(ans)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
A = list(map(int, input().split()))
MOD = 10**9 + 7

# case1. N == K のときは答えは確定
if N == K:
    ans = 1
    for i in range(K):
        ans *= A[i]
        ans %= MOD
    print(ans)
# case2. A が全部負の数のとき、 Kが奇数ならば残念ながら答えは負になるのでなるべく絶対値を小さくする
elif max(A) < 0 and K % 2 == 1:
    ans = 1
    A = sorted(A, reverse=True)
    for i in range(K):
        ans *= A[i]
        ans %= MOD
    print(ans)
else:
    # case3. それ以外は正にできるので、絶対値の降順にソートする
    A = sorted(A, key=lambda x: abs(x), reverse=True)
    # とりあえず大きい順に K 個掛けてみる
    ans = 1
    for i in range(K):
        ans *= A[i]
    if ans >= 0:
        ans %= MOD
        print(ans)
    else:
        # 取り除く対象を old, 取り込む対象を new とする
        # case3-1. 負数を取り除いて非負数を取り込む
        # まずは new を探す
        j = K
        while j < N or A[j] < 0:
            j += 1
        if j < N:
            new1 = A[j]
        # old を探す
        l = K-1
        while A[l] >= 0:
            l -= 1
        old1 = A[l]
        if new1 and old1:
            cand1 = ans * pow(old1, MOD-2, MOD) * new1
        # case3-2. 非負数を取り除いて負数を取り込む
        # まずは new を探す
        j = K
        while j < N or A[j] >= 0:
            j += 1
        if j < N:
            new2 = A[j]
        # old を探す
        l = K-1
        while A[l] < 0:
            l -= 1
        old2 = A[l]
        if new2 and old2:
            cand2 = ans * pow(old2, MOD-2, MOD) * new2

        if cand1 and cand2:
            print(cand1, cand2)
            ans = max(cand1, cand2) % MOD
        elif cand1:
            ans = cand1 % MOD
        else:
            ans = cand2 % MOD
        print(ans)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
A = list(map(int, input().split()))
MOD = 10**9 + 7

# case1. N == K のときは答えは確定
if N == K:
    ans = 1
    for i in range(K):
        ans *= A[i]
        ans %= MOD
    print(ans)
# case2. A が全部負の数のとき、 Kが奇数ならば残念ながら答えは負になるのでなるべく絶対値を小さくする
elif max(A) < 0 and K % 2 == 1:
    ans = 1
    A = sorted(A, reverse=True)
    for i in range(K):
        ans *= A[i]
        ans %= MOD
    print(ans)
else:
    # case3. それ以外は正にできるので、絶対値の降順にソートする
    A_ = sorted(A, key=lambda x: abs(x), reverse=True)
    # とりあえず大きい順に K 個掛けてみる
    ans = 1
    for i in range(K):
        ans *= A[i]
    if ans >= 0:
        ans %= MOD
        print(ans)
    else:
        # 取り除く対象を old, 取り込む対象を new とする
        # case3-1. 負数を取り除いて非負数を取り込む
        # まずは new を探す
        j = K
        while j < N or A[j] < 0:
            j += 1
        if j < N:
            new1 = A[j]
        # old を探す
        l = K-1
        while A[l] >= 0:
            l -= 1
        old1 = A[l]
        if new1 and old1:
            cand1 = ans * pow(old1, MOD-2, MOD) * new1
        # case3-2. 非負数を取り除いて負数を取り込む
        # まずは new を探す
        j = K
        while j < N or A[j] >= 0:
            j += 1
        if j < N:
            new2 = A[j]
        # old を探す
        l = K-1
        while A[l] < 0:
            l -= 1
        old2 = A[l]
        if new2 and old2:
            cand2 = ans * pow(old2, MOD-2, MOD) * new2

        if cand1 and cand2:
            print(cand1, cand2)
            ans = max(cand1, cand2) % MOD
        elif cand1:
            ans = cand1 % MOD
        else:
            ans = cand2 % MOD
        print(ans)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
A = list(map(int, input().split()))
MOD = 10**9 + 7

# case1. N == K のときは答えは確定
if N == K:
    ans = 1
    for i in range(K):
        ans *= A[i]
        ans %= MOD
    print(ans)
# case2. A が全部負の数のとき、 Kが奇数ならば残念ながら答えは負になるのでなるべく絶対値を小さくする
elif max(A) < 0 and K%2 == 1:
    ans = 1
    A = sorted(A, reverse=True)
    for i in range(K):
        ans *= A[i]
        ans %= MOD
    print(ans)
else:
    # case3. それ以外は正にできるので、絶対値の降順にソートする
    A_ = sorted(A, key=lambda x: abs(x), reverse=True)
    # とりあえず大きい順に K-1 個掛けてみる
    ans = 1
    for i in range(K-1):
        ans *= A[i]
    if ans >= 0:
        ans %= MOD
        print(ans)
    else:
        # 取り除く対象を old, 取り込む対象を new とする
        # case3-1. 負数を取り除いて非負数を取り込む
        # まずは new を探す
        j = K
        while j < N or A[j] < 0:
            j += 1
        if j < N:
            new1 = A[j]
        # old を探す
        l = K-1
        while A[l] >= 0:
            l -= 1
        old1 = A[l]
        if new1 and old1:
            cand1 = ans * pow(old1, MOD-2, MOD) * new1
        # case3-2. 非負数を取り除いて負数を取り込む
        # まずは new を探す
        j = K
        while j < N or A[j] >= 0:
            j += 1
        if j < N:
            new2 = A[j]
        # old を探す
        l = K-1
        while A[l] < 0:
            l -= 1
        old2 = A[l]
        if new2 and old2:
            cand2 = ans * pow(old2, MOD-2, MOD) * new2

        if cand1 and cand2:
            ans = max(cand1, cand2) % MOD
        elif cand1:
            ans = cand1 % MOD
        else:
            ans = cand2 % MOD
        print(ans)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
A = list(map(int, input().split()))
MOD = 10**9 + 7

# case1. N == K のときは答えは確定
if N == K:
    ans = 1
    for i in range(K):
        ans *= A[i]
        ans %= MOD
    print(ans)
    exit()

# case2. A が全部負の数のとき、 Kが奇数ならば残念ながら答えは負になるのでなるべく絶対値を小さくする
if max(A) < 0 and K%2 == 1:
    ans = 1
    A = sorted(A, reverse=True)
    for i in range(K):
        ans *= A[i]
        ans %= MOD
    print(ans)
    exit()

# case3. それ以外は、負の数は基本的に2つセットで選ばれることに注意
# まずは候補テーブルを作成
A = sorted(A, reverse=True)
A_cost1 = []  # positive 要素
A_cost2 = []  # negative 要素を大きい順に 2 つセットにしたもの
tmp = 1
for a in A:
    if a >= 0:
        A_cost1.append(a)
    else:
        if tmp != 1:
            A_cost2.append(tmp*a)
            tmp = 1
        else:
            tmp *= a
A_cost1 = sorted(A_cost1, reverse=True)
A_cost2 = sorted(A_cost2, reverse=True)

cost = 0
ans = 1
while (cost < K):
    if K - cost >= 2:
        # まだ cost2 から入れる余地がある
        if A_cost2[0] >= A_cost1[0]:
            res = A_cost2.pop(0)
            ans *= res
            ans %= MOD
            cost += 2
        else:
            res = A_cost1.pop(0)
            ans *= res
            ans %= MOD
            cost += 1
    else:
        res = A_cost1.pop(0)
        ans *= res
        ans %= MOD
        cost += 1

print(ans)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
A = list(map(int, input().split()))
MOD = 10**9 + 7

# case1. N == K のときは答えは確定
if N == K:
    ans = 1
    for i in range(K):
        ans *= A[i]
        ans %= MOD
    print(ans)
    exit()

A_abs = [abs(a) for a in A]
A_is_negative = [a < 0 for a in A]

# A_abs の大きい順にソート
A_abs, A_is_negative, A = zip(*sorted(zip(A_abs, A_is_negative, A), key=lambda x: x[0], reverse=True))

# case2. A が全部負の数のとき、 Kが奇数ならば残念ながら答えは負になるのでなるべく絶対値を小さくする
if len(A_is_negative) == N and K%2 == 1:
    ans = 1
    for i in range(K):
        ans *= A[-(i+1)]
        ans %= MOD
    print(ans)
    exit()

# case3. それ以外は、負の数は基本的に2つセットで選ばれることに注意
# まずは候補テーブルを作成
A_cost1 = []  # positive 要素
A_cost2 = []  # negative 要素を大きい順に 2 つセットにしたもの
tmp = 1
for a in A:
    if a >= 0:
        A_cost1.append(a)
    else:
        if tmp != 1:
            A_cost2.append(tmp*a)
            tmp = 1
        else:
            tmp *= a

cost = 0
ans = 1
while (cost < K):
    if K - cost >= 2:
        # まだ cost2 から入れる余地がある
        if A_cost2[0] >= A_cost1[0]:
            res = A_cost2.pop(0)
            ans *= res
            ans %= MOD
            cost += 2
        else:
            res = A_cost1.pop(0)
            ans *= res
            ans %= MOD
            cost += 1
    else:
        res = A_cost1.pop(0)
        ans *= res
        ans %= MOD
        cost += 1

print(ans)