Problem p02774 (python)

Fixed submission:
# 小さい方からk番目がx
# x以下がk個以上となる最小のx
# 負の数のところでバグってそう
from bisect import bisect_left, bisect_right

import sys

input = sys.stdin.buffer.readline

n, k = map(int, input().split())
A = list(map(int, input().split()))
A.sort()
minusA = bisect_left(A, 0)
mA = A[:minusA]

zeroA = bisect_right(A, 0) - bisect_left(A, 0)
plusA = n - minusA - zeroA
pA = A[-plusA:]
minus = minusA * plusA
plus = plusA * (plusA - 1) // 2 + minusA * (minusA - 1) // 2
zero = zeroA * (zeroA - 1) // 2 + (n - zeroA) * zeroA


# x以下の数を数える
def calc(x):
    mmA = mA[::-1]
    res = 0
    if x > 0:
        lim = plusA - 1
        for i in range(plusA):
            while lim >= 0:
                if pA[i] * pA[lim] > x:
                    lim -= 1
                else:
                    break
            res += lim + 1
            if lim >= i:
                res -= 1
        lim = minusA - 1
        mmA = mA[::-1]
        for i in range(minusA):
            while lim >= 0:
                if mmA[i] * mmA[lim] > x:
                    lim -= 1
                else:
                    break
            res += lim + 1
            if lim >= i:
                res -= 1
        res //= 2
        res += minus + zero
    elif x == 0:
        res += minus + zero
    else:
        # 正の数、左絶対値小さい
        # 負の数、左絶対値大きい
        lim = minusA - 1
        for i in range(plusA):
            while lim >= 0:
                if pA[i] * mmA[lim] <= x:
                    lim -= 1
                else:
                    break
            res += minusA - (lim + 1)
    return res


ng = -(10 ** 20)
ok = 10 ** 20
while ok - ng > 1:
    mid = (ok + ng) // 2
    if calc(mid) >= k:
        ok = mid
    else:
        ng = mid

print(ok)

Buggy submission (verdict: Runtime Error):
# 片方、しゃくとりにして高速化
from bisect import bisect_left, bisect_right

n, k = map(int, input().split())
A = list(map(int, input().split()))
A.sort()

zero_num = bisect_right(A, 0) - bisect_left(A, 0)
minus_num = bisect_left(A, 0)
plus_num = n - (zero_num + minus_num)
AP = A[zero_num + minus_num :]
AM = A[:minus_num]
RAM = AM[::-1]
plus = plus_num * (plus_num - 1) // 2 + minus_num * (minus_num - 1) // 2
zero = zero_num * (zero_num - 1) // 2 + zero_num * (n - zero_num)
minus = plus_num * minus_num

# ↑ここまでで正負の積の数をそれぞれ求めておく。


def count(x):
    cnt = 0
    if x > 0:
        right = plus_num
        for left in range(plus_num):
            seki = AP[left] * AP[right - 1]
            while seki > x and right > 0:
                right -= 1
                seki = AP[left] * AP[right - 1]
            cnt += right
        for i in range(plus_num):
            if AP[i] ** 2 <= x:
                cnt -= 1
        right = minus_num
        for left in range(minus_num):
            seki = RAM[left] * RAM[right - 1]
            while seki > x and right > 0:
                right -= 1
                seki = RAM[left] * RAM[right - 1]
            cnt += right
        for i in range(minus_num):
            if RAM[i] ** 2 <= x:
                cnt -= 1
        cnt //= 2
        cnt += minus + zero
    elif x < 0:
        right = minus_num
        for left in range(plus_num):
            seki = AP[left] * AM[right - 1]
            while seki > x and right > 0:
                right -= 1
                seki = AP[left] * RAM[right - 1]
            cnt += right
    else:
        cnt = minus + zero
    return cnt


ok = 10 ** 20
ng = -(10 ** 20)
while ok - ng > 1:
    mid = (ok + ng) // 2
    if count(mid) >= k:
        ok = mid
    else:
        ng = mid
print(ok)

Buggy submission (verdict: Runtime Error):
# 片方、しゃくとりにして高速化
from bisect import bisect_left, bisect_right

n, k = map(int, input().split())
A = list(map(int, input().split()))
A.sort()

zero_num = bisect_right(A, 0) - bisect_left(A, 0)
minus_num = bisect_left(A, 0)
plus_num = n - (zero_num + minus_num)
AP = A[zero_num + minus_num :]
AM = A[:minus_num]
RAM = AM[::-1]
plus = plus_num * (plus_num - 1) // 2 + minus_num * (minus_num - 1) // 2
zero = zero_num * (zero_num - 1) // 2 + zero_num * (n - zero_num)
minus = plus_num * minus_num

# ↑ここまでで正負の積の数をそれぞれ求めておく。


def count(x):
    cnt = 0
    if x > 0:
        right = plus_num  # 半開区間で考えるので[0,1)から始める
        for left in range(plus_num):
            seki = AP[left] * AP[right - 1]
            while seki > x and right > 0:
                right -= 1
                seki = AP[left] * AP[right - 1]
            cnt += right
        for i in range(plus_num):
            if AP[i] ** 2 <= x:
                cnt -= 1
        right = minus_num
        for left in range(minus_num):
            seki = RAM[left] * RAM[right - 1]
            while seki > x and right > 0:
                right -= 1
                seki = RAM[left] * RAM[right - 1]
            cnt += right
        for i in range(minus_num):
            if RAM[i] ** 2 <= x:
                cnt -= 1
        cnt //= 2
        cnt += minus + zero
    elif x < 0:
        right = 1  # 半開区間で考えるので[0,1)から始める
        for left in range(plus_num):
            seki = AP[left] * RAM[right - 1]
            while seki > x and right < minus_num:
                right += 1
                seki = AP[left] * RAM[right - 1]
            cnt += minus_num - right + 1
    else:
        cnt = minus + zero
    return cnt


ok = 10 ** 20
ng = -(10 ** 20)
while ok - ng > 1:
    mid = (ok + ng) // 2
    if count(mid) >= k:
        ok = mid
    else:
        ng = mid
print(ok)