Problem p02774 (python)

Fixed submission:
from bisect import bisect_left, bisect_right
import sys
input = sys.stdin.readline

n, k = map(int, input().split())
a = list(map(int, input().split()))
a.sort()

a_minus = a_zero = a_plus = 0
for i in range(n):
    if a[i] > 0:
        a_plus += 1
    elif a[i] == 0:
        a_zero += 1
    else:
        a_minus += 1

aa_minus = a_minus * a_plus
aa_zero = a_zero * (n - a_zero) + a_zero * (a_zero - 1) // 2


def check_minus(x):
    # k番目の数がx(<0)以下であるか
    # x以下の数字がk個以上あるか
    cnt = 0
    j = a_minus + a_zero
    for i in range(a_minus):
        # a[i] < 0
        while j < n and a[i] * a[j] > x:
            j += 1
        cnt += n - j

    return cnt >= k


def check_plus(x):
    # k番目の数がx(>0)以下であるか
    # x以下の数字がk個以上あるか

    def func(b, m):
        # 昇順ソート済みの数列b (長さm)
        # 掛け算してx以下となる数を求める
        c = 0
        j = m - 1
        for i in range(m):
            while j > 0 and b[i] * b[j] > x:
                j -= 1
            c += max(j - i, 0)
        return c

    cnt = aa_minus + aa_zero
    cnt += func(a[-a_plus:], a_plus)
    cnt += func(list(reversed(a[:a_minus])), a_minus)
    return cnt >= k


if k <= aa_minus:
    # k番目の数は負の数
    lb = a[0] * a[-1] - 1  # False
    ub = 0  # True
    while ub - lb > 1:
        mid = (ub + lb) // 2
        if check_minus(mid):
            ub = mid
        else:
            lb = mid

    ans = ub
elif k <= aa_minus + aa_zero:
    # k番目の数は0
    ans = 0
else:
    # k番目の数はプラス
    lb = 0  # False
    ub = max(a[0]**2, a[-1]**2) + 1  # True
    while ub - lb > 1:
        mid = (ub + lb) // 2
        if check_plus(mid):
            ub = mid
        else:
            lb = mid

    ans = ub

print(ans)

Buggy submission (verdict: Runtime Error):
from bisect import bisect_left, bisect_right
n, k = map(int, input().split())
a = list(map(int, input().split()))
a.sort()

n_m = bisect_left(a, 0)  # a[i] < 0 の数
n_0 = bisect_right(a, 0) - n_m  # a[i] = 0 の数
n_p = n - n_m - n_0  # a[i] > 0 の数

nn_m = n_m * n_p  # 積が負になるペア数
nn_0 = n_0 * (n - n_0) + n_0 * (n_0 - 1) // 2  # 積が0になるペア数
nn_p = n * (n - 1) // 2 - nn_m - nn_0  # 積が正になるペア数

if k <= nn_m:
    # k番目の値は負

    def check_minus(x):
        # k番目の値がx(<0)以下

        # a[i] * a[j] <= x

        # a[i]   * a[j] <= x  : a[j] >= abs(x) / abs(a[i])
        # a[i+1] * a[j] <= x  : a[j] >= abs(x) / abs(a[i+1])
        #     a[i]  <=     a[i+1] < 0
        # abs(a[i]) >= abs(a[i+1])
        # a: [-------00000+++++++]
        #     i             jjjjj
        #      i             jjjj
        #       i              jj

        # a[i] * a[j] <= -7
        # n_m = 4
        # n_0 = 3
        # a = [-5, -3, -2, -1, 0, 0, 0, 1, 2, 3, 4, 5]
        #       i                          j (j  j  j)
        #           i                         j (j  j)
        #               i                        j
        #                   i                          j
        # a[i] * a[j] = -5
        #  -2  *   4  = -6

        cnt = 0
        j = n_m + n_0
        for i in range(n_m):
            while a[i] * a[j] > x and j < n:
                j += 1
            # a[i] * a[j] <= x
            # [j, n)
            cnt += n - j

        return cnt >= k

    # a[i] * a[j] <= x
    lb = -10**18  # False
    ub = -1  # True
    while ub - lb > 1:
        mid = (ub + lb) // 2
        if check_minus(mid):
            ub = mid
        else:
            lb = mid
    ans = ub
elif k <= nn_m + nn_0:
    # k番目の値は0
    ans = 0
else:
    # k番目の値は正
    a_p = a[n_m + n_0:]
    a_m = [-a[i] for i in reversed(range(n_m))]

    def check_plus(x):
        # k番目の値がx(>=0)以下
        def func(b):
            # b [++++++++++++++]
            #    ijjjjjjjjjjj
            #     ijjjjjjjj
            #      ijjjj
            #       ij

            # b[i] * b[j] <= x
            # b[i]   * b[j] <= x : b[j]  <= x / b[i]
            # b[i+1] * b[j] <= x : b[j]  <= x / b[i+1]

            c = 0
            j = len(b) - 1
            for i in range(len(b)):
                if i >= j:
                    break
                # i < j
                while i < j and b[i] * b[j] > x:
                    j -= 1
                # b[i] * b[j] <= x
                # (i, j]
                c += j - i
            return c

        cnt = nn_m + nn_0
        cnt += func(a_p)
        cnt += func(a_m)
        return cnt >= k

    # a[i] * a[j] <= x
    lb = 0  # False
    ub = 10**18 + 1  # True
    while ub - lb > 1:
        mid = (ub + lb) // 2
        if check_plus(mid):
            ub = mid
        else:
            lb = mid
    ans = ub
print(ans)

Buggy submission (verdict: Runtime Error):
from bisect import bisect_left, bisect_right
n, k = map(int, input().split())
a = tuple(map(int, input().split()))
a.sort()


def check(x):
    # k番目の数がx以下であるか
    # x以下の数字がk個以上あるか
    cnt = 0
    for i, ai in enumerate(a):
        if ai > 0:
            idx = bisect_right(a, x // ai)
            cnt += max(idx - i - 1, 0)
        elif ai < 0:
            idx = bisect_left(a, -(x // abs(ai)))
            # cnt += n - idx
            cnt += n - max(idx, i + 1)
        else:
            # ai == 0
            if x >= 0:
                cnt += n - i - 1

        if cnt >= k:
            return True
    # print('check', x, cnt)
    return cnt >= k


lb = -10**18  # False
ub = +10**18  # True
while ub - lb > 1:
    mid = (ub + lb) // 2
    # print(lb, ub)
    if check(mid):
        ub = mid
    else:
        lb = mid

print(ub)