Problem p02774 (python)

Fixed submission:
import bisect
import numpy as np

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

#N,K = 9,9
#A = [-3,-2,-2,-2,0,0,1,2,2]

def simple():
    B = sorted([A[i] * A[j] for i in range(N) for j in range(i+1,N)])
    return B[K-1]

def solve():
    A.sort()

    iplus = bisect.bisect_right(A,0)
    izeros = bisect.bisect_left(A,0)

    pluss = np.array(A[iplus:]).astype(np.int64)
    zeros = A[izeros:iplus]
    minuss = np.array(A[:izeros]).astype(np.int64)

    nplus = len(pluss)
    nzero = len(zeros)
    nminus = len(minuss)

    npminus = nplus * nminus
    npzeros = nzero * (nplus + nminus) + nzero * (nzero-1)//2

    if npminus < K <= npminus + npzeros:
        return 0
    elif K <= npminus:
        mn = A[-1] * A[0] - 1
        mx = 0
        while mx-mn > 1:
            mid = (mn+mx)//2
            xs = -(-mid//minuss)
            cnt = nplus * nminus
            cnt -= np.sum(np.searchsorted(pluss,xs,side="left"))
            if cnt >= K:
                mx = mid
            else:
                mn = mid
        return mx
    else:
        nK = K - npminus - npzeros
        pluss2 = (-minuss)[::-1]
        mn = -1
        mx = max(A[0]**2, A[-1]**2)
        pluss_p2 = pluss**2
        pluss2_p2 = pluss2**2
        while mx-mn > 1:
            mid = (mn+mx)//2
            cnt = 0
            xs = mid//pluss
            cnt += np.sum(np.searchsorted(pluss,xs,side='right'))
            cnt -= np.count_nonzero(pluss_p2 <= mid)
            xs2 = mid//pluss2
            cnt += np.sum(np.searchsorted(pluss2,xs2,side='right'))
            cnt -= np.count_nonzero(pluss2_p2 <= mid)
            cnt = cnt //2
            #print(mn,mid,mx,cnt)
            if nK <= cnt:
                mx = mid
            else:
                mn = mid
        return mx

print(solve())
#for K in range(1,N*(N-1)//2+1):
#    print(solve(),simple())

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

A.sort()

iplus = bisect.bisect_right(A,0)
izeros = bisect.bisect_left(A,0)

pluss = A[iplus:]
zeros = A[izeros:iplus]
minuss = A[:izeros]

nplus = len(pluss)
nzero = len(zeros)
nminus = len(minuss)

npminus = nplus * nminus
npzeros = nzero * (nplus + nminus) + nzero * (nzero-1)//2

if npminus < K <= npminus + npzeros:
    print(0)
elif K <= npminus:
    mn = minuss[0] * pluss[-1] - 1
    mx = minuss[-1] * pluss[0]
    while mx-mn > 1:
        mid = (mn+mx)//2
        cnt = 0
        for m in minuss:
            x = -(-mid//m)
            k = bisect.bisect_left(pluss,x)
            #print(m,x,pluss,k)
            cnt += (nplus-k)
            if cnt >= K:
                break
        if cnt >= K:
            mx = mid
        else:
            mn = mid
    print(mx)
else:
    K -= npminus + npzeros
    pluss2 = [-m for m in minuss][::-1]
    mn = pluss[0] * pluss[1] - 1
    mx = pluss[-1] * pluss[-2]
    if len(pluss2) > 1:
        mn = min(mn, pluss2[0] * pluss2[1] - 1)
        mx = max(mx, pluss2[-1] * pluss2[-2])

    while mx-mn > 1:
        mid = (mn+mx)//2
        cnt = 0
        for m in pluss:
            x = mid//m
            k = bisect.bisect_right(pluss,x)
            cnt += k - (m**2 <= mid)
        for m in pluss2:
            x = mid//m
            k = bisect.bisect_left(pluss2,x)
            cnt += k - (m**2 <= mid)
        cnt = cnt //2
        #print(mn,mid,mx,cnt)
        if cnt >= K:
            mx = mid
        else:
            mn = mid
    print(mx)