Problem p02774 (python)

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

A = [A[i] for i in range(N) if A[i]!=0]
c = (N+len(A)-1)*(N-len(A))
N = len(A)
A.sort()
B = A[::-1]

def count(m):
    res = 0
    if m>=0:
        res += c
        j = N
        for i in range(N):
            if A[i]>0:
                break
            while j>0:
                if A[i]*A[j-1]<=m:
                    j -= 1
                else:
                    break
            res += N-j
        if A[i]>0:
            j = N-1
            for i in range(i,N):
                while j>=0:
                    if A[i]*A[j]>m:
                        j -= 1
                    else:
                        break
                res += j+1
    else:
        j = 0
        for i in range(N):
            if B[i]<0:
                break
            while j<N:
                if B[i]*B[j]>m:
                    j += 1
                else:
                    break
            res += N-j
        if B[i]<0:
            j = -1
            for i in range(i,N):
                while j<N-1:
                    if B[i]*B[j+1]<=m:
                        j += 1
                    else:
                        break
                res += j+1
    for i in range(N):
        if A[i]*A[i]<=m:
            res -= 1
    res //= 2
    return res

ng = -10**18
ok = 10**18
m = (ng+ok)//2
if N==0:
    ans = 0
else:
    for n in range(100):
        if count(m)<K:
            ng = m
        else:
            ok = m
        m = (ng+ok)//2
    ans = int(ok)
print(ans)

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

A = [A[i] for i in range(N) if A[i]!=0]
c = (N+len(A)-1)*(N-len(A))
N = len(A)
A.sort()
B = A[::-1]

def count(m):
    res = 0
    if m>=0:
        res += c
        j = N
        for i in range(N):
            if A[i]>0:
                break
            while j>0:
                if A[i]*A[j-1]<=m:
                    j -= 1
                else:
                    break
            res += N-j
        if A[i]>0:
            j = N-1
            for i in range(i,N):
                while j>=0:
                    if A[i]*A[j]>m:
                        j -= 1
                    else:
                        break
                res += j+1
    else:
        j = N
        for i in range(N):
            if B[i]<0:
                break
            while j>0:
                if B[i]*B[j-1]<=m:
                    j -= 1
                else:
                    break
            res += N-j
        if B[i]<0:
            j = -1
            for i in range(i,N):
                while j<N-1:
                    if B[i]*B[j+1]<=m:
                        j += 1
                    else:
                        break
                res += j+1
    for i in range(N):
        if A[i]*A[i]<=m:
            res -= 1
    res //= 2
    return res

ng = -10**18
ok = 10**18
m = (ng+ok)//2
for n in range(100):
    if count(m)<K:
        ng = m
    else:
        ok = m
    m = (ng+ok)//2
ans = int(ok)
print(ans)

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

def count(m):
    res = 0
    for i in range(N):
        a = A[i]
        if a==0:
            if m>=0:
                res += N
        elif a>0:
            r = m/a
            j = bisect(A,r)
            res += j
        else:
            r = m/a
            j = bisect_left(A,r)
            res += N-j
        if a**2<=m:
            res -= 1
    res //= 2
    return res,cand

def count1(m):
    res = 0
    cand = -10**19
    for i in range(N):
        a = A[i]
        if a==0:
            if m>=0:
                res += N
        elif a>0:
            r = m/a
            j = bisect(A,r)
            res += j
            if j!= 0:
                cand = max(cand,a*A[j-1])
        else:
            r = m/a
            j = bisect_left(A,r)
            res += N-j
            if j != N:
                cand = max(cand,a*A[j])
        if a**2<=m:
            res -= 1
    res //= 2
    return res,cand


l = -10**18
r = 10**18
m = (l+r)/2
for i in range(60):
    k = count(m)
    if k < K:
        l = m
    else:
        r = m
    m = (l+r)/2
k,ans = count1(r)
print(ans)