Problem p02616 (python)

Fixed submission:
import sys; input = sys.stdin.buffer.readline
sys.setrecursionlimit(10**7)
from collections import defaultdict
con = 10 ** 9 + 7; INF = float("inf")

def getlist():
	return list(map(int, input().split()))

#処理内容
def main():
	N, K = getlist()
	A = getlist()
	B = []
	for i in range(N):
		val = abs(A[i])
		if A[i] > 0:
			B.append([val , 1])
		elif A[i] < 0:
			B.append([val, -1])
		else:
			B.append([val, 0])
	B = sorted(B)

	if N == K:
		ans = 1
		for i in range(N):
			ans *= A[i]
			ans %= con
		print(ans)
		return

	M = 0; p = 0; m = 0; Aplus = []; Aminus = []; zerocnt = 0
	for i in A:
		if i == 0:
			zerocnt += 1
		elif i > 0:
			Aplus.append(i)
			M += 1; p += 1
		else:
			Aminus.append(i)
			M += 1; m += 1
	Aplus.sort()
	Aminus.sort()

	# 0が絶対に入る場合
	if K > M:
		print(0)
		return
	# マイナスになるから絶対ゼロを入れる場合
	if zerocnt > 0 and M == m and K % 2 == 1:
		print(0)
		return
	# マイナスのみの時の処理
	if m == N:
		if K % 2 == 0:
			ans = 1
			for i in range(K):
				ans *= Aminus[i]
				ans %= con
			print(ans)
			return
		else:
			ans = 1
			for i in range(K):
				ans *= Aminus[-1 - i]
				ans %= con
			print(ans)
			return

	# マイナスと0のとき
	if m == M:
		if K % 2 == 0:
			ans = 1
			for i in range(K):
				ans *= Aminus[i]
				ans %= con
			print(ans)
			return
		else:
			print(0)
			return

	pindPre = -100; pindAfter = INF
	mindPre = -100;	mindAfter = INF
	# print(B)

	for i in range(N - K):
		if B[i][1] == 1:
			pindPre = i
		elif B[i][1] == -1:
			mindPre = i

	for i in range(N - 1, N - K - 1, -1):
		if B[i][1] == 1:
			pindAfter = i
		elif B[i][1] == -1:
			mindAfter = i

	cnt = 0
	ans = 1
	for i in range(K):
		ans *= B[-1 - i][0]
		ans %= con
		if B[-1 - i][1] == -1:
			cnt += 1

	if cnt % 2 == 0:
		print(ans)
		return
	else:
		# ans *= -1
		ans %= con
		if mindPre >= 0 and pindPre >= 0 and pindAfter <= N and mindAfter <= N:
			if B[mindPre][0] * B[mindAfter][0] <= B[pindPre][0] * B[pindAfter][0]:
				ans *= B[pindPre][0] * pow(B[mindAfter][0], con - 2, con)
				ans %= con
			else:
				ans *= B[mindPre][0] * pow(B[pindAfter][0], con - 2, con)
				ans %= con
			print(ans)
			return

		elif pindPre >= 0 and mindAfter <= N:
			ans *= B[pindPre][0] * pow(B[mindAfter][0], con - 2, con)
			ans %= con
			print(ans)
			return
		elif mindPre >= 0 and pindAfter <= N:
			ans *= B[mindPre][0] * pow(B[pindAfter][0], con - 2, con)
			ans %= con
			print(ans)
			return
		else:
			print(0)
			return


if __name__ == '__main__':
	main()

Buggy submission (verdict: Runtime Error):
import sys; input = sys.stdin.buffer.readline
sys.setrecursionlimit(10**7)
from collections import defaultdict
con = 10 ** 9 + 7; INF = float("inf")

def getlist():
	return list(map(int, input().split()))

#処理内容
def main():
	N, K = getlist()
	A = getlist()
	B = []
	for i in range(N):
		val = abs(A[i])
		if A[i] > 0:
			B.append([val , 1])
		elif A[i] < 0:
			B.append([val, -1])
		else:
			B.append([val, 0])
	B = sorted(B)

	M = 0; p = 0; m = 0; Aplus = []; Aminus = []; zerocnt = 0
	for i in A:
		if i == 0:
			zerocnt += 1
		elif i > 0:
			Aplus.append(i)
			M += 1; p += 1
		else:
			Aminus.append(i)
			M += 1; m += 1
	Aplus.sort()
	Aminus.sort()

	# 0が絶対に入る場合
	if K > M:
		print(0)
		return
	# マイナスになるから絶対ゼロを入れる場合
	if zerocnt > 0 and M == m and K % 2 == 1:
		print(0)
		return
	# マイナスのみの時の処理
	if m == N:
		if K % 2 == 0:
			ans = 1
			for i in range(K):
				ans *= Aminus[i]
				ans %= con
			print(ans)
			return
		else:
			ans = 1
			for i in range(K):
				ans *= Aminus[-1 - i]
				ans %= con
			print(ans)
			return

	# マイナスと0のとき
	if m == M:
		if K % 2 == 0:
			ans = 1
			for i in range(K):
				ans *= Aminus[i]
				ans %= con
			print(ans)
			return
		else:
			print(0)
			return

	pindPre = -100; pindAfter = INF
	mindPre = -100;	mindAfter = INF
	# print(B)

	for i in range(N - K):
		if B[i][1] == 1:
			pindPre = i
		elif B[i][1] == -1:
			mindPre = i

	for i in range(N - 1, N - K - 1, -1):
		if B[i][1] == 1:
			pindAfter = i
		elif B[i][1] == -1:
			mindAfter = i

	cnt = 0
	ans = 1
	for i in range(K):
		ans *= B[-1 - i][0]
		ans %= con
		if B[-1 - i][1] == -1:
			cnt += 1

	if cnt % 2 == 0:
		print(ans)
		return
	else:
		# ans *= -1
		ans %= con
		if mindPre >= 0 and pindPre >= 0 and pindAfter <= N and mindAfter <= N:		
			if B[mindPre][0] * B[mindAfter][0] <= B[pindPre][0] * B[pindAfter][0]:
				ans *= B[pindPre][0] * pow(B[mindAfter][0], con - 2, con)
				ans %= con
			else:
				ans *= B[mindPre][0] * pow(B[pindAfter][0], con - 2, con)
				ans %= con
			print(ans)
			return

		elif pindPre >= 0 and mindAfter <= N:
			ans *= B[pindPre][0] * pow(B[mindAfter][0], con - 2, con)
			ans %= con
			print(ans)
			return
		elif mindPre >= 0 and pindAfter <= N:
			ans *= B[mindPre][0] * pow(B[pindAfter][0], con - 2, con)
			ans %= con
			print(ans)
			return
		else:
			print0(0)
			return


if __name__ == '__main__':
	main()