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 sys
from bisect import bisect_left
import numpy as np


def test(a, K):
    from itertools import combinations
    answer = []
    for i, j in combinations(a, 2):
        answer.append(i * j)
    answer.sort()
    print(answer)
    print(answer[K - 1])


def main():
    read = sys.stdin.read

    N, K, *A = map(int, read().split())

    test(A, K)
    A.sort()
    m = bisect_left(A, 0)
    z = bisect_left(A, 1) - m
    p = N - m - z

    M = m * p
    Z = z * (N - z) + z * (z - 1) // 2
    negative = np.array(A[:m], np.int64)
    positive = np.array(A[-p:], np.int64)

    if K <= M:
        negative = -negative[::-1]
        left = 1
        right = -A[0] * A[-1] + 1
        while left + 1 < right:
            mid = (left + right) // 2
            cnt = (p - np.searchsorted(positive, np.ceil(mid / negative).astype(int), side='left')).sum()
            if cnt < K:
                right = mid
            else:
                left = mid
        print(-left)
    elif K <= M + Z:
        print(0)
    else:
        K -= M + Z
        negative = np.abs(negative[::-1])
        left = 0
        right = max(A[0] ** 2, A[-1] ** 2)
        while left + 1 < right:
            mid = (left + right) // 2
            a = mid // negative
            b = np.searchsorted(negative, a, side='right').sum() - (mid >= negative ** 2).sum()
            b //= 2

            d = mid // positive
            e = np.searchsorted(positive, d, side='right').sum() - (mid >= positive ** 2).sum()
            e //= 2

            if b + e < K:
                left = mid
            else:
                right = mid

        print(right)


if __name__ == '__main__':
    main()


import sys
from bisect import bisect_left
import numpy as np
 
 
def test(a, K):
    from itertools import combinations
    answer = []
    for i, j in combinations(a, 2):
        answer.append(i * j)
    answer.sort()
    print(answer[K - 1])
 
 
def main():
    read = sys.stdin.read
 
    N, K, *A = map(int, read().split())
 
    if N <= 10:
        test(A, K)
        exit()
    A.sort()
    m = bisect_left(A, 0)
    z = bisect_left(A, 1) - m
    p = N - m - z
 
    M = m * p
    Z = z * (N - z) + z * (z - 1) // 2
    negative = np.array(A[:m], np.int64)
    positive = np.array(A[-p:], np.int64)
    negative = -negative[::-1]
    

import sys
from bisect import bisect_left
import numpy as np


def test(a, K):
    from itertools import combinations
    answer = []
    for i, j in combinations(a, 2):
        answer.append(i * j)
    answer.sort()
    print(answer)
    print(answer[K - 1])


def main():
    read = sys.stdin.read

    N, K, *A = map(int, read().split())

    test(A, K)
    A.sort()
    m = bisect_left(A, 0)
    z = bisect_left(A, 1) - m
    p = N - m - z

    M = m * p
    Z = z * (N - z) + z * (z - 1) // 2
    minuses = np.array(A[:m], np.int64)
    pluses = np.array(A[-p:], np.int64)

    if K <= M:
        left = A[0] * A[-1] - 1
        right = -1
        while left + 1 < right:
            
            mid = (left + right) // 2
            cnt = (p - np.searchsorted(pluses, np.ceil(mid / minuses).astype(int), side='left')).sum()
            if cnt < K:
                left = mid
            else:
                right = mid

        print(right)
    elif K <= M + Z:
        print(0)
    else:

        K -= M+Z
        left = 0
        right = max(minuses[0] ** 2, pluses[-1] ** 2)
        while left + 1 < right:
            mid = (left + right) // 2
            a = mid // minuses
            b = (m - np.searchsorted(minuses, a, side='left')).sum() - (mid >= minuses ** 2).sum()
            b //= 2

            d = mid // pluses
            e = np.searchsorted(pluses, d, side='right').sum() - (mid >= pluses ** 2).sum()
            e //= 2

            if b + e < K:
                left = mid
            else:
                right = mid

        print(right)


if __name__ == '__main__':
    main()


import sys
import numpy as np


def test(a, K):
    from itertools import combinations
    answer = []
    for i, j in combinations(a, 2):
        answer.append(i * j)
    answer.sort()
    print(answer[K - 1])


def main():
    read = sys.stdin.read

    N, K, *A = map(int, read().split())

    if N <= 10:
        test(A, K)
        exit()
    A.sort()
    A = np.array(A, np.int64)
    negative = A[A < 0]
    positive = A[A > 0]
    m = len(negative)
    p = len(positive)
    z = N - m - p

    M = m * p
    Z = z * (N - z) + z * (z - 1) // 2
    negative = -negative[::-1]

import sys
from bisect import bisect_left
import numpy as np


def test(a, K):
    from itertools import combinations
    answer = []
    for i, j in combinations(a, 2):
        answer.append(i * j)
    answer.sort()
    print(answer[K - 1])


def main():
    read = sys.stdin.read

    N, K, *A = map(int, read().split())

    if N <= 10:
        test(A, K)
        exit()
    A.sort()
    m = bisect_left(A, 0)
    z = bisect_left(A, 1) - m
    p = N - m - z

    M = m * p
    Z = z * (N - z) + z * (z - 1) // 2
    if m > 0:
    	negative = np.array(A[:m], np.int64)
    else:
      	negative = np.array([], np.int64)
    if p > 0:
    	positive = np.array(A[-p:], np.int64)
    else:
      	positive = np.array([], np.int64)
    negative = -negative[::-1]
    

    if K <= M:
        left = 1
        right = -A[0] * A[-1] + 1
        while left + 1 < right:
            mid = (left + right) // 2
            cnt = (p - np.searchsorted(positive, (mid + negative - 1) // negative, side='left')).sum()
            if cnt < K:
                right = mid
            else:
                left = mid
        print(-left)
    elif K <= M + Z:
        print(0)
    else:
        K -= M + Z
        left = 0
        right = max(A[0] * A[1], A[-1] * A[-2])
        positive2 = positive ** 2
        negative2 = negative ** 2
        while left + 1 < right:
            mid = (left + right) // 2
            a = np.searchsorted(negative, mid // negative, side='right').sum() - (mid >= negative2).sum()
            b = np.searchsorted(positive, mid // positive, side='right').sum() - (mid >= positive2).sum()

            if (a + b) // 2 < K:
                left = mid
            else:
                right = mid

        print(right)


if __name__ == '__main__':
    main()


Verdicts: Wrong Answer

Wrong Answer

Runtime Error

Wrong Answer

Accepted