ARC 037 C - 億マス計算 - BEN2のブログ

二分探索を 2 回使う

問題 : ARC 037 C - 億マス計算

$N \times N$ マス計算で得られる $N^{2}$ 個の値を昇順に並べ、小さい方から $K$ 番目の値を出力する問題。
$N$ は最大 $3 \times 10^{4}$ なので、 $N^{2}$ 個のマスを 1 個ずつ埋めていく方法は TLE となる。

以下、典型的な解法として覚えておく。
小さい方から $K$ 番目の値を $X$ とする。このとき、次のことが言える。

$X-1$ 以下であるマスは $K$ 個未満しか存在せず、 $X$ 以下であるマスは $K$ 個以上存在する

よって、「 $x$ 以下であるマスが $K$ 個以上存在する」を満たす最小の $x$ (これが $X$ ) を求めればよい。

$x$ が小さいほど $x$ 以下のマスの個数は少なく、 $x$ が大きいほど $x$ 以下のマスの個数は多くなるため、実装では、二分探索によって $X$ を求める。

では、 $x$ 以下であるマスが $K$ 個以上存在するかどうかの判定はどうやるか。これも二分探索を用いる。
あらかじめ $b_{j}$ を昇順にソートして、 $N \times N$ マス表を、どの行も右へいくほど値が大きくなるようにしておく。

f:id:BEN2suzuka:20200504205854p:plain

以下の方法で 1 行ずつ $x$ 以下のマスの個数を数えて、合計の個数を求めるとよい。

$i$ 行目について、 $\displaystyle b_{j} \leq \left\lfloor \frac{x}{a_{i}} \right\rfloor$ を満たす $j$ の個数を求める

$b_{j}$ のうち、要素が $\displaystyle \left\lfloor \frac{x}{a_{i}} \right\rfloor$ 以下である $j$ の個数は二分探索で求められる。

#include <bits/stdc++.h>
using namespace std;

int N, K;
vector<long long> A, B;

// x 以下であるマスが K 個以上あるかどうか判定
bool func(long long x) {
  int res = 0;
  // 1 行ずつ x 以下であるマスを数える
  for (int i = 0; i < N; i++) {
    long long tmp = x / A.at(i);
    // b_j <= [x / a_i] を満たす j の個数を加算
    res += upper_bound(B.begin(), B.end(), tmp) - B.begin();
  }
  return res >= K;
}

int main() {
  cin >> N >> K;
  A.resize(N);
  B.resize(N);
  for (int i = 0; i < N; i++) cin >> A.at(i);
  for (int i = 0; i < N; i++) cin >> B.at(i);
  sort(B.begin(), B.end());
  // x 以下であるマスが K 個以上あるような最小の x を求める
  long long left = 0;
  long long right = 1LL << 60;
  while (right - left > 1) {
    long long mid = (left + right) / 2;
    if (func(mid)) right = mid;
    else left = mid;
  }
  cout << right << endl;
}