Problem p02774 (python)

Fixed submission:
# 単純な組み合わせだと計算量がO(n**2)となってしまうので工夫が必要
# 小さい順にK番目の値を求めよ
# ⇄ 判定条件「x以下のものがK個以上ある」を満たす最小のxを求めよ
# ⇄ ・判定条件を満たす最小のx：2分探索 (-10**18 ~ 10**18)
# 　 ・「数列から２個選んだ積がx以下のものが何個あるかを数えて、K個以下か判定せよ」：片方を固定して2分探索で数えていく
import sys
input = sys.stdin.readline

N, K = [int(x) for x in input().split()]
A = [int(x) for x in input().split()]
A = sorted(A)
ok, ng = 10**18, -10**18

# 判定条件を満たす最小のx：2分探索 (-10**18 ~ 10**18)
while ok - ng > 1:
    mid = (ok + ng) // 2
    # 数列から２個選んだ積がx以下のものが何個あるかを数える
    s = 0
    t = 0
    for i in range(N):
        if A[i] > 0:
            # 固定したA[i]が正の場合、積の数列
            # A[i]*A[0], A[i]*A[1], ... , A[i]*A[n-1]
            # つまり、{A[i]*A[j]} j=0~n-1 という数列は単調増加
            # j=JOKの時が「x以下」という条件を満たす最大値とすると、条件を満たす個数はJNG個である
            # 図解すると、
            # x以下(OK), x以下(OK), ..., x以下(OK), xより大きい(NG),...,xより大きい(NG)
            # <------------ JNG個 --------------><----------- N - JNG個 ----------->
            # 二分探索でこの境界値Jを求めにいく
            jok, jng = -1, N
            while jng - jok > 1:
                jmid = (jok + jng) // 2
                if A[i] * A[jmid] <= mid:
                    jok = jmid
                else:
                    jng = jmid
            s += jok + 1
        elif A[i] < 0:
            # 固定したA[i]が負の場合、積の数列
            # A[i]*A[0], A[i]*A[1], ... , A[i]*A[n-1]
            # は単調減少。
            # j=JOKの時が、「x以下」という条件を満たす最小のJだとすると、条件を満たす個数はJ+1個である
            # 図解すると、
            # xより大きい(NG), xより大きい(NG), ..., xより大きい(NG), x以下(OK),...,x以下(OK)
            # <--------------------- JOK個 ---------------------><----- N-JOK個 ------->
            jng, jok = -1, N
            while jok - jng > 1:
                jmid = (jok + jng) // 2
                if A[i] * A[jmid] <= mid:
                    jok = jmid
                else:
                    jng = jmid
            s += (N - jok)
        elif mid >= 0:
            s += N
        if A[i] * A[i] <= mid:
            t += 1
    # 求めたsは条件を満たす順列の個数なので、組み合わせの個数にするために/2する
    total = (s - t) // 2
    # それがK個以上あるかを判定する
    if total >= K:
        ok = mid
    else:
        ng = mid
print(ok)

Buggy submission (verdict: Runtime Error):
# 単純な組み合わせだと計算量がO(n**2)となってしまうので工夫が必要
# 小さい順にK番目の値を求めよ
# ⇄ 判定条件「x以下のものがK個以上ある」を満たす最小のxを求めよ
# ⇄ ・判定条件を満たす最小のx：2分探索 (-10**18 ~ 10**18)
# 　 ・「数列から２個選んだ積がx以下のものが何個あるかを数えて、K個以下か判定せよ」：片方を固定して2分探索で数えていく
import sys
input = sys.stdin.readline

N, K = [int(x) for x in input().split()]
A = [int(x) for x in input().split()]
A = sorted(A)
ok, ng = 10**18, -10**18

# 判定条件を満たす最小のx：2分探索 (-10**18 ~ 10**18)
while ok - ng > 1:
    mid = (ok + ng) // 2
    # 数列から２個選んだ積がx以下のものが何個あるかを数える
    s = 0
    t = 0
    for i in range(N):
        if A[i] > 0:
            # 固定したA[i]が正の場合、積の数列
            # A[i]*A[0], A[i]*A[1], ... , A[i]*A[n-1]
            # つまり、{A[i]*A[j]} j=0~n-1 という数列は単調増加
            # j=JOKの時が「x以下」という条件を満たす最大値とすると、条件を満たす個数はJNG個である
            # 図解すると、
            # x以下(OK), x以下(OK), ..., x以下(OK), xより大きい(NG),...,xより大きい(NG)
            # <------------ JNG個 --------------><----------- N - JNG個 ----------->
            # 二分探索でこの境界値Jを求めにいく
            jok, jng = -1, N
            while jng - jok > 1:
                jmid = (jok + jng) // 2
                if A[i] * A[jmid] <= x:
                    jok = jmid
                else:
                    jng = jmid
            s += jok + 1
        elif A[i] < 0:
            # 固定したA[i]が負の場合、積の数列
            # A[i]*A[0], A[i]*A[1], ... , A[i]*A[n-1]
            # は単調減少。
            # j=JOKの時が、「x以下」という条件を満たす最小のJだとすると、条件を満たす個数はJ+1個である
            # 図解すると、
            # xより大きい(NG), xより大きい(NG), ..., xより大きい(NG), x以下(OK),...,x以下(OK)
            # <--------------------- JOK個 ---------------------><----- N-JOK個 ------->
            jng, jok = -1, N
            while jok - jng > 1:
                jmid = (jok + jng) // 2
                if A[i] * A[jmid] <= mid:
                    jok = jmid
                else:
                    jng = jmid
            s += (N - jok)
        elif mid >= 0:
            s += N
        if A[i] * A[i] <= mid:
            t += 1
    # 求めたsは条件を満たす順列の個数なので、組み合わせの個数にするために/2する
    total = (s - t) // 2
    # それがK個以上あるかを判定する
    if total >= K:
        ok = mid
    else:
        ng = mid
print(ok)