AtCoder が提供しているABC（AtCoder Beginner Contest）364 D問題をC++とPythonで解いてみました。ABC364は、2024年7月27日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
D問題 K-th Nearest（Difficulty : 1136）
問題の詳細は、リンク先をご覧ください。
解の二分探索を使いますが、判定条件も二分探索で判断します。AtCoder Problems による Difficulty は 1136 でした。
解答案
$K$ 番目に近い点を求めるのではなく、「距離 $d$ 以内の点が $K$ 個以上あるか」を判定します（解の二分探索と呼んでいます）。
解の二分探索について、めぐる式と呼ばれる二分探索の実装を参考にしました。
判定の考え方は、ABC360D問題（解説記事）と同じで、$L \leqq a[i] \leqq R$ となる配列の要素 $a_i$ の個数を求めます。
C++ プログラム例（ABC364D）
数直線上の $N$ 個の点をソートしておきます。この問題では、判定関数を分けずに埋め込みました（28ー29行目）。
- 23行目から35行目までは、解の二分探索の定番コードです。
- upper_bound から lower_bound を引いて、閉区間に含まれる要素数を求めています（28行目）。
以下が、C++プログラムです。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
#define INF (1LL << 60)
int main()
{
int n, q;
cin >> n >> q;
vector<ll> a(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
}
sort(a.begin(), a.end());
for (int i = 0; i < q; ++i) {
ll b;
int k;
cin >> b >> k;
ll ng = -1;
ll ok = 1000000000;
while (abs(ok - ng) > 1) {
ll mid = (ok + ng) / 2;
int c = upper_bound(a.begin(), a.end(), b + mid) - lower_bound(a.begin(), a.end(), b - mid);
if (c >= k) {
ok = mid;
} else {
ng = mid;
}
}
cout << ok << endl;
}
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC364D）
Python版も基本的な考え方は同じです。二分探索は、bisect を使いました。bisect_left が lower_bound に、bisect_right が upper_bound に相当します。以下がプログラムです。
"""AtCoder Beginner Contest 364 D"""
import bisect
n, q = map(int, input().split())
a = list(map(int, input().split()))
a.sort()
for i in range(q):
b, k = map(int, input().split())
ng = -1
ok = 10 ** 9
while abs(ok - ng) > 1:
mid = (ok + ng) // 2
c = bisect.bisect_right(a, b + mid) - \
bisect.bisect_left(a, b - mid)
if c >= k:
ok = mid
else:
ng = mid
print(ok)こちらも「AC」と判定されました。
最後に
最終的に定番コードが多いプログラムとなりましたが、コンテストでは1時間以上かかりました。以下で時間を使っていました。
- upper_bound と lower_bound の使い方の混乱
- 解の二分探索の変数 ng の初期値を 0 としていた。このため、距離が同じ点を別にカウントして加えていた。
まぁ解けて良かったですが、次からはもう少し早く解きたいです。
引き続き ABC の問題を紹介していきます。
