D問題 K-th Nearest
この記事では、AtCoderの問題を解説します。
「二分探索」を応用する問題でした。今の私にぴったりな問題なので、記事としてメモ化しておきます。Pythonでの解き方を知りたい方にもピッタリです。
問題
公式による解説はこちら
https://atcoder.jp/contests/abc364/editorial/10549
解き方のポイント
- 「k番目に○○な△△を求める」という問題には、二分探索が有効。
- fj(x)という関数を考える。これは、a1からa2のうち、Bjとの距離がx以下であるようなものの数。
- fj(x) >= kjを満たす最小のxを求めればいい。
- xの範囲を絞り込むために二分探索が使える。
注意
二分探索を行う際、以下の条件が満たされることを保証する必要がある。
- ng は常に不適切な値（条件を満たさない値）を表す。
- ok は常に適切な値（条件を満たす値）を表す。
具体的には、探索中の距離xが ng と ok の間に位置しているとき、条件を満たすかどうかを確認する。
初期状態で ng = -1 とする理由は、距離xの最小値が 0 だから。
Pythonのコード
import bisect
# 入力を受け取る
n, q = map(int, input().split())
a = list(map(int, input().split()))
# リストをソート
a.sort()
# クエリ処理
for _ in range(q):
b, k = map(int, input().split())
def f(x):
# [b-x, b+x] の範囲内にある点の数がk以上かどうかをチェックする
lb = bisect.bisect_left(a, b - x)
ub = bisect.bisect_right(a, b + x)
return ub - lb >= k
ng, ok = -1, int(2e8) + 10
while ok - ng > 1:
mid = (ok + ng) // 2
if f(mid):
ok = mid
else:
ng = mid
print(ok)