配列のなかから指定された値に対して K 番目に近い値を取得する方法を考えます。
Contents
D – K-th Nearest
数直線上に N + Q 個の点 A_1, … , A_N, B_1, …, B_Q があり、点 A_i の座標は a_i、点 B_j の座標は b_j です。
j = 1, 2, …, Q それぞれについて、以下の問題に答えてください。
点 A_1, A_2, …, A_N のうち点 B_j との距離が k_j 番目に近い点を X としたとき、点 X と点 B_j との距離を求めよ。
入力されるデータ
N Q
a_1 a_2 … a_N
b_1 k_1
b 2 k_2
…
b_Q k_Qただし、
1 ≦ N, Q ≦ 10^5
– 10^8 ≦ a_i, b_j ≦ 10^8
1 ≦ k_j ≦ N
これはある値が指定されたときに配列のなかからK 番目にその値に近い値を取得する問題とみることができます。最大で 10^5 回のクエリが飛んでくるので処理を高速化する方法を考えなければなりません。
二分探索法で探す
このような問題は二分探索が有効です。まず配列 A をソートしましょう。そして配列の要素のうち X – m 以上 X + m 以下の個数が K になる m の最小値を探します。
二分探索法で探すのでBinarySearchクラスを定義します。コンストラクタには配列 A を渡します。
1 2 3 4 5 6 7 8 9 10 11 12 |
class BinarySearch { long[] A = new long[1]; int N = 0; public BinarySearch(long[] vs) { N = vs.Length; A = (long[])vs.Clone(); A = A.OrderBy(_ => _).ToArray(); } } |
前提として指定された要素以上の値が現れる最小のインデックスを取得するLowerBoundメソッド、指定された要素より大きい値が現れる最小のインデックスを取得するUpperBoundメソッド、指定された要素が存在するかを確認するExistメソッドをを定義します。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
class BinarySearch { public int LowerBound(long target) { if (A[0] >= target) return 0; int left = 0; int right = A.Length; while (right - left > 1) { int mid = (left + right) / 2; if (target <= A[mid]) right = mid; else left = mid; } return right; } public int UpperBound(long target) { if (A[0] > target) return 0; int left = 0; int right = A.Length; while (right - left > 1) { int mid = (left + right) / 2; if (target < A[mid]) right = mid; else left = mid; } return right; } public bool Exist(long target) { if (A[0] > target) return false; if (A[0] == target) return true; int left = 0; int right = A.Length; while (right - left > 1) { int mid = (left + right) / 2; if (target == A[mid]) return true; if (target < A[mid]) right = mid; else left = mid; } return false; } } |
二分探索法で target に K 番目に近い値を取得します。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
class BinarySearch { // K 番目に近い値を取得する public long GetNearestK(long target, int k) { // b - mid 以上 b + mid 以下 の要素の個数が K 以上か? bool f(long b, long mid) { int lb = LowerBound(b - mid); int ub = UpperBound(b + mid); return ub - lb >= k; } long left = -1; // right の取りうる最大値を計算する long max = Math.Max(A[N - 1], target); long min = Math.Min(A[0], target); long right = max - min + 1; while (right - left > 1) { long mid = (right + left) / 2; if (f(target, mid)) right = mid; else left = mid; } // K 番目に近い値は(target - right)と(target + right)のうち存在する側 if (Exist(target + right)) return target + right; if (Exist(target - right)) return target - right; throw new Exception("ここは実行されないはず"); } } |
定義したクラスを用いて解を求めます。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 |
class Program { static void Main() { int N, Q; { int[] vs = Console.ReadLine().Split().Select(_ => int.Parse(_)).ToArray(); N = vs[0]; Q = vs[1]; } long[] A = Console.ReadLine().Split().Select(_ => long.Parse(_)).ToArray(); A = A.OrderBy(_ => _).ToArray(); BinarySearch bs = new BinarySearch(A); for (int i = 0; i < Q; i++) { int[] vs = Console.ReadLine().Split().Select(_ => int.Parse(_)).ToArray(); int b = vs[0]; int k = vs[1]; long v = bs.GetNearestK(b, k); long d = Math.Abs(v - b); Console.WriteLine(d); } } } |