yutasの競技プログラミング勉強帖

競技プログラミングの問題についての解説記事を主に書いています。

D - Reindeer and Sleigh / ユニークビジョンプログラミングコンテスト2023 クリスマス (AtCoder Beginner Contest 334)

ABC334 AtCoder C++ AtCoder Beginner Contest 400点 二分探索 競技プログラミング

問題

n 台のソリがあり、各ソリには 1, 2, \cdots , n の番号が付けられている。 また、ソリ i を引くのに必要なトナカイは {r}_{i} 匹であり、各トナカイが引けるソリの数は高々 1 台である。

以下の形式のクエリが q 個与えられるので、そのそれぞれに答えよ。

  • 整数 x が与えられる。 x 匹のトナカイがいるときの、ソリを引ける最大数を求めよ。

入力

最初の1行目に整数 n, q が与えられる。

次の1行に整数 {r}_{1}, {r}_{2}, \cdots , {r}_{n} が与えられる。

その後 q 行に渡って、整数 x が与えられる。

条件

  • 実行時間制限: 2s
  • メモリ制限: 1024MB
  • 1 \leq n, q \leq 2 \times {10}^{5}
  • 1 \leq {r}_{i} \leq {10}^{9}
  • 1 \leq x \leq 2 \times {10}^{14}

出力

q 行に渡り、各クエリの解答を出力すること。

解法

引けるソリの数をできるだけ大きくする問題なので、なるべくトナカイの必要数が少ないものから順番にソリを選ぶ必要があります。 そのため、 {r}_{1}, \cdots , {r}_{n} を昇順にソートしたものとして以下考えます。

ここで、 \mathrm{sum}_{i} を、 1 台目から i 台目までのソリを引くのに必要なトナカイの総数として定めます。 すなわち、 \begin{align} \mathrm{sum}_{i} = {r}_{1} + \cdots + {r}_{i} \end{align} として定めます。 （ただし便宜上、 \mathrm{sum}_{0} = 0 とします。）

このようにして作成された \mathrm{sum} を用いて、ある整数 x に対して \begin{align} \mathrm{sum}_{k} \leq x < \mathrm{sum}_{k + 1} \end{align} を満たす整数 k が、求める答えとなります。

この整数 k を求めるために、 \mathrm{sum}_{1} から順に見ていくと、全体で O(nq) の計算量がかかってしまいます。 これでは、 n, q \leq 2 \times {10}^{5} であるので、実行時間制限に間に合いません。

そこで、二分探索を用いて上記の k の値を求めることを考えます。 この方法を用いることによって、全体で O(q \log n) の計算量に削減することができるので、実行時間制限に間に合います。

よって、これが解法となります。 全体の計算量としては、 {r}_{1}, \cdots , {r}_{n} をソートする部分で O(n \log n) かかり、クエリで O(q \log n) かかるので、 O ((n + q) \log n) かかることになります。

ソースコード

まずは二分探索で k の値を求める部分について、以下のような solve() 関数を実装しました。

int n;
ll sum[int(2e5 + 5)];

int solve(ll x) {
  if (sum[n] <= x) { // 与えられた x が sum[n] 以上であれば、 n を返す
    return n;
  }

  int l = 0, r = n; // 二分探索の端となる値をそれぞれ設定する
  while (r - l > 1) {
    int mid = (l + r) / 2; // l, r の中間の値を比較対象とする
    if (sum[mid] <= x) { // もし、sum[mid] <= x ならば、左端を mid にする
      l = mid;
    } else { // そうでないならば、右端を mid とする
      r = mid;
    }
  }
  return l; // 最終的に sum[l] <= x < sum[r] となるので、 l を返す
}

この solve() 関数がある状態で、 main() 関数の中にクエリの答えを出力する部分を実装しました。

int n, q;
ll r[int(2e5 + 5)], sum[int(2e5 + 5)];

int main() {
  // 値の入力
  cin >> n >> q;
  for (int i = 1; i <= n; i++) {
    cin >> r[i];
  }

  // r の中身を昇順にソートする
  sort(r + 1, r + n + 1);

  // sum の値を計算する
  for (int i = 1; i <= n; i++) {
    sum[i] = sum[i - 1] + r[i];
  }

  // q 個のクエリに答える
  for (int i = 1; i <= q; i++) {
    ll x;
    cin >> x; // x の値の入力
    cout << solve(x) << endl; // クエリの答えの出力
  }

  return 0;
}

感想

かなりシンプルな二分探索の問題だと思いました。 初めて二分探索を実装する方への練習問題として良さそうに思います。