Problem: We have N integers A_1, A_2, ..., A_N.  There are \frac{N(N-1)}{2} ways to choose two of them and form a pair. If we compute the product of each of those pairs and sort the results in ascending order, what will be the K-th number in that list?

Submission attempts (in order):
import numpy as np
N,K=map(int,input().split())
A=list(map(int,input().split()))
A=sorted(A)
#B=list(reversed(A))
A=np.array(A)
Aneg=A[A<0]
Azero=A[A==0]
Apos=A[A>0]

n_neg=len(Aneg)*len(Apos)
n_zero=len(Azero)*(len(Azero)-1)//2+len(Azero)*(len(Aneg)+len(Apos))
n_pos=(len(Apos)*(len(Apos)-1)+len(Aneg)*(len(Aneg)-1))//2

if K <=n_neg:
    n=K-n_neg
    B=sorted(Aneg)[:n+1]
    C=sorted(Apos)[:n+1]
    ans=[]
    for i in range(min(len(B),n-1)):
        for j in range(min(len(C),n-1)):
            ans.append(B[i]*C[j])
    ans=sorted(ans)
    print(ans[n-1])
    
elif K>n_neg+n_zero:
    n=K-(n_neg+n_zero)
    B=sorted(Aneg)
    B=list(reversed(B))
    C=sorted(Apos)
    ans=[]
    if len(B)>2:
        for i in range(len(B)-1):
            for j in range(i+1,len(B)):
                ans.append(B[i]*B[j])
    if len(C)>2:
#         k=min(len(C),n-1)
        for i in range(len(C)-1):
            for j in range(i+1,len(C)):
                ans.append(C[i]*C[j])
    ans=sorted(ans)
    print(ans[n-1])
    
else:
    print(0)

from ABC155_D2 import *



import sys
from io import StringIO
import unittest

class TestClass(unittest.TestCase):
    def assertIO(self, input, output):
        stdout, stdin = sys.stdout, sys.stdin
        sys.stdout, sys.stdin = StringIO(), StringIO(input)
        resolve()
        sys.stdout.seek(0)
        out = sys.stdout.read()[:-1]
        sys.stdout, sys.stdin = stdout, stdin
        self.assertEqual(out, output)
    def test_入力例_1(self):
        input = """4 3
3 3 -4 -2"""
        output = """-6"""
        self.assertIO(input, output)
    def test_入力例_2(self):
        input = """10 40
5 4 3 2 -1 0 0 0 0 0"""
        output = """6"""
        self.assertIO(input, output)
    def test_入力例_3(self):
        input = """30 413
-170202098 -268409015 537203564 983211703 21608710 -443999067 -937727165 -97596546 -372334013 398994917 -972141167 798607104 -949068442 -959948616 37909651 0 886627544 -20098238 0 -948955241 0 -214720580 277222296 -18897162 834475626 0 -425610555 110117526 663621752 0"""
        output = """448283280358331064"""
        self.assertIO(input, output)

if __name__ == "__main__":
    unittest.main()

import numpy as np
N,K=map(int,input().split())
A=list(map(int,input().split()))

A=np.array(A)

A_pos=np.sort(A[A>0])
A_neg=np.sort(A[A<0])
A_neg2=-A_neg[::-1]
n_pos=len(A_pos)
n_neg=len(A_neg)
n_zero=N-n_pos-n_neg

def position_in_neg(x):
    y=n_pos-np.searchsorted(A_pos,x//A_neg,side='right')
    return y.sum()

def position_in_pos(x):
    z=np.searchsorted(A_pos,-(-x//A_pos),side='left')
    w=np.searchsorted(A_neg2,-(-x//A_neg2),side='left')
    tmp=len(A_pos[A_pos**2<x])+len(A_neg2[A_neg2**2<x])
    return (z.sum()+w.sum()-tmp)//2

def position(x):
    if x<0:
        return position_in_neg(x)
    elif x==0:
        return n_pos*n_neg
    else:
        return position_in_pos(x)+n_pos*n_neg+(n_pos+n_neg)*n_zero+n_zero*(n_zero-1)//2

l=-pow(10,18)
r=pow(10,18)
while(r-l>1):
    mid=(l+r)//2
    if position(mid)<K:
        l=mid
    else:
        r=mid
print(l)

Verdicts: Runtime Error

Runtime Error

Accepted