Problem p02774 (python)

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

pos = []
neg = []
zero = 0
for a in A:
    if a > 0:
        pos.append(a)
    elif a == 0:
        zero += 1
    else:
        neg.append(a)
        
len_pos = len(pos)
len_neg = len(neg)

cnt_pos = len_pos * (len_pos - 1) // 2 + len_neg * (len_neg - 1) // 2
cnt_neg = len_pos * len_neg
cnt_zero = zero * (zero - 1) // 2 + zero * (len_pos + len_neg)
#print(cnt_pos, cnt_neg, cnt_zero, cnt_pos + cnt_neg + cnt_zero, N * (N - 1) // 2)

if K <= cnt_neg:
    pos.sort()
    neg.sort()
    ok = 10 ** 19
    ng = -10 ** 19
    while abs(ok - ng) > 1:
        mid = (ok + ng) // 2
        pos_i, neg_i = len_pos - 1, len_neg - 1
        cnt = 0
        while pos_i >= 0 and neg_i >= 0:
            if pos[pos_i] * neg[neg_i] <= mid:
                cnt += neg_i + 1
                pos_i -= 1
            else:
                if neg_i:
                    neg_i -= 1
                else:
                    pos_i -= 1
        if cnt >= K:
            ok = mid
        else:
            ng = mid
    print(ok)

elif K <= cnt_neg + cnt_zero:
    print(0)
    
else:
    K -= cnt_neg + cnt_zero
    pos.sort()
    neg = [abs(n) for n in neg]
    neg.sort()

    ok = 10 ** 19
    ng = -10 ** 19
    while abs(ok - ng) > 1:
        mid = (ok + ng) // 2
        left, right = 0, len_pos - 1
        cnt = 0
        while left <= len_pos - 1 and right >= 0:
            if pos[left] * pos[right] <= mid:
                cnt += right + 1
                if left <= right:
                    cnt -= 1
                left += 1
            else:
                if right:
                    right -= 1
                else:
                    left += 1

        left, right = 0, len_neg - 1
        while left <= len_neg - 1 and right >= 0:
            if neg[left] * neg[right] <= mid:
                cnt += right + 1
                if left <= right:
                    cnt -= 1
                left += 1
            else:
                if right:
                    right -= 1
                else:
                    left += 1
        cnt //= 2
        
        if cnt >= K:
            ok = mid
        else:
            ng = mid
    print(ok)

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.buffer.readline

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

plus, minus, zero = [], [], 0
for a in A:
    if a > 0:
        plus.append(a)
    elif a < 0:
        minus.append(a)
    else:
        zero += 1
        
len_plus = len(plus)
len_minus = len(minus)

plus_cnt = len_plus * (len_plus - 1) // 2 + len_minus * (len_minus - 1) // 2
minus_cnt = len_plus * len_minus
zero_cnt = N * (N - 1) // 2 - plus_cnt - minus_cnt
print(plus_cnt, minus_cnt, zero_cnt)

# 小さい順に K 番目に来る数字を求める
if K <= minus_cnt:
    # plus(4, 2, 1), minus(-1, -2, -3)
    # 4, -2, 2, -3
    plus.sort(reverse = True)
    minus.sort(reverse = True)
    plus.append(-10 ** 18)
    minus.append(10 ** 18)

    ok, ng = 10 ** 19, -10 ** 19
    while abs(ok - ng) > 1:
        mid = (ok + ng) // 2
        l, r = 0, 0
        v = plus[l] * minus[r]
        cnt = 0
        while l < len_plus and r < len_minus:
            if v <= mid:
                cnt += len_minus - r
                l += 1
            else:
                r += 1
            v = plus[l] * minus[r]
        if cnt >= K:
            ok = mid
        else:
            ng = mid
    print(ok)
elif minus_cnt < K <= minus_cnt + zero_cnt:
    print(0)
else:
    1
    """
    # plus(3, 2, 1), minus(5, 4, 2)
    minus = [abs(m) for m in minus]
    plus.sort(reverse = True)
    minus.sort(reverse = True)
    plus.append(-10 ** 18)
    minus.append(-10 ** 18)

    ok, ng = 10 ** 19, -10 ** 19
    while abs(ok - ng) > 1:
        mid = (ok + ng) // 2
        l, r = 0, 0
        v = plus[l] * plus[r]
        cnt = 0
        while l < len_plus:
            if v <= mid:
                cnt += (len_minus - r)
                l += 1
            else:
                r += 1
            v = plus[l] * minus[r]
        if cnt >= K:
            ok = mid
        else:
    """
            ng = mid
    print(ok)