ABC 155 D - Pairs

問題

長さ $n$ の配列 $A$ から $2$ つ選んでできるペアの積のうち $k$ 番目に小さいものを求めよ。（$n \leq 2\times 10 ^ 5$）

atcoder.jp

解説

$k$ 番目を求めるには、「$x$ より小さいものの個数が $k$ 未満になるような最大の $x$」 を求めればよい。
積が $x$ より小さいペアの個数は

1. $A _ i < 0$ なる各 $i$ に対して $A _ i \times A _ j > x$ なる $j$ の個数
2. $x > 0$ の場合のみ、$A _ j = 0$ なる $j$ の個数 $\times n$
3. $A _ i > 0$ なる各 $i$ に対して $A _ i \times A _ j < x$ なる $j$ の個数

を合計して、$A _ i ^ 2 < x$ なる $i$ の個数を引いてから $2$ で割ったものである。
あとは「$x$ より小さいものの個数が $k$ 未満になるような最大の $x$」を二分探索で求めればよい。

実装

import numpy as np

def count_pair(x):
    cnt = 0
    cnt += (n-np.searchsorted(A, x//A_nega, side='right')).sum()
    if x > 0:
        cnt += len(A_zero)*n
    cnt += np.searchsorted(A, -(-x//A_posi), side='left').sum()
    cnt = (cnt-(A**2 < x).sum())//2
    return cnt

n, k = map(int, input().split())
* A, = map(int, input().split())
A = np.array(A)
A = np.sort(A)
A_posi = A[A > 0]
A_zero = A[A == 0]
A_nega = A[A < 0]
inf = 10**18+1
l, r = -inf, inf
while r - l > 1:
    c = (l + r) // 2
    if count_pair(c) < k:
        l = c
    else:
        r = c
print(l)

atcoder.jp

類題

$A$ がすべて正のバージョン。 atcoder.jp