Official

G - K番目の要素/K-th element Editorial by Nyaan


まずは計算量を無視して考えると、各 \(i\) \((1 \leq i \leq N)\) に対して愚直に数列の中身を列挙すれば \(\mathrm{O}(\sum_i A_i)\) 程度で解くことができますが TLE してしまいます。

この問題のように \(K\) 番目の要素を求める問題では、二分探索 を利用したアルゴリズムが有効になることが多いです。
答えを \(a\) とおくと、 \(K\) 以上の要素の個数と \(a\) の間に以下の関係が成り立つことがわかります。

  • \(x \lt a\) のとき、 \(x\) 以下の要素は \(K\) 個未満
  • \(x \geq a\) のとき、 \(x\) 以下の要素は \(K\) 個以上

よって、各 \(x\) に対して 「 \(x\) 以下の要素は \(K\) 個以上あるか?」という判定問題を十分高速に解くことができれば、二分探索を利用して \(\mathrm{O}(\log M)\) 回判定問題を解くことでこの問題を解くことができます。 ( \(M\) は数列の要素の最大値)

判定問題を高速に解く方法を考えます。
長さ \(A_i\), 初項 \(B_i\), 公差 \(C_i\) の数列に \(x\) 以下の要素が何個含まれているかを考えます。これは初項と \(x\) の大小で場合分けをすると次の関係が導かれます。

  • \(x \lt B_i\) のとき \(0\)
  • \(x \geq B_i\) のとき \(\left\lfloor \frac{x - B_i}{C_i} \right\rfloor + 1\)

よって各数列に対して \(x\) 以下の要素の個数が \(\mathrm{O}(1)\) で求まるので、判定問題は \(N\) 個の数列に対して 1 つずつ判定を行うことで \(\mathrm{O}(N)\) で解くことができました。

以上より、この問題を全体で \(\mathrm{O}(N \log M)\) で解くことができました。以下に C++ による解答例を載せます。

#include <iostream>
#include <vector>
using namespace std;
 
int main() {
  long long N, K;
  cin >> N >> K;
  vector<long long> A(N), B(N), C(N);
  for (int i = 0; i < N; i++) cin >> A[i] >> B[i] >> C[i];
  // x以下の要素がK個以上か?で二分探索
  long long ng = 0, ok = 1LL << 60;
  while (ng + 1 < ok) {
    long long x = (ng + ok) / 2;
    long long count = 0;
    for (int i = 0; i < N; i++) {
      if (x < B[i]) continue;
      long long num = (x - B[i]) / C[i] + 1;
      count += min(A[i], num);
    }
    (K <= count ? ok : ng) = x;
  }
  long long ans = ok;
  cout << ans << "\n";
}

posted:
last update: