AOJ 2527: MLE

そもそも乱数列全部見るだけで厳しそう、ということは乱数列の生成の仕方の性質で数学ゲーするのかな、と判断し、分からないので解説を見ました。 xorshiftは早いので間に合う、とのことでした。

solution

愚直にやるとMLEするし、sortの部分でTLEもする。 しっかりsortすべきなのは$k$番目付近だけなので、ここだけ取り出してsortすればよい。$O(N)$。

空間全体を適当に分割し、それぞれの区間に入る乱数の数を数える。 これにより、$k$番目の数がどの区間に入るかが分かるので、その区間の数だけ取り出してsortすればよい。 対象は乱数であるので、どの区間にもほぼ均等に数が入ると考えてよく、間に合う。

注意として、

• $x_0 = 0$のときは周期が$1$に潰れるので、答えはでるが$n$が大きいとTLEる
  • なのでxorshiftの周期は$2^{64}$でなく$2^{64}-1$
• a[uint64_t(x) >> 56] += 1;とかすると符号の順序を見ていないためにWAる

implementation

#include <iostream>
#include <vector>
#include <algorithm>
#define repeat(i,n) for (int i = 0; (i) < (n); ++(i))
#define whole(f,x,...) ([&](decltype((x)) y) { return (f)(begin(y), end(y), ## __VA_ARGS__); })(x)
using namespace std;
template <typename F>
void xorshift64(int n, uint64_t x, F f) {
    repeat (i,n) {
        f(x);
        x ^= x << 13;
        x ^= x >> 7;
        x ^= x << 17;
    }
}
int main() {
    int n, k, x0; cin >> n >> k >> x0; -- k;
    if (x0 == 0) {
        cout << 0 << endl;
        return 0;
    }
    vector<int> cnt(256);
    xorshift64(n, x0, [&](uint64_t x) {
        cnt[int8_t(x >> 56) + 128] += 1;
    });
    int acc = 0;
    int j = 0;
    for (; acc + cnt[j] <= k; ++ j) {
        acc += cnt[j];
    }
    vector<uint64_t> a;
    xorshift64(n, x0, [&](uint64_t x) {
        if (int8_t(x >> 56) + 128 == j) a.push_back(x);
    });
    whole(sort, a);
    cout << int64_t(a[k - acc]) << endl;
    return 0;
}