docker_space/frontier_cs_4/local_test.cpp

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

// ---- Matrix generation ----
struct TestCase {
    int n;
    long long k;
    vector<vector<long long>> A; // 1-indexed
    long long answer;
};

mt19937_64 rng(42);

TestCase gen_additive(int n, long long k) {
    // a[i][j] = i + j
    TestCase tc;
    tc.n = n; tc.k = k;
    tc.A.assign(n+1, vector<long long>(n+1, 0));
    vector<long long> all;
    for (int i = 1; i <= n; i++)
        for (int j = 1; j <= n; j++) {
            tc.A[i][j] = i + j;
            all.push_back(i + j);
        }
    sort(all.begin(), all.end());
    tc.answer = all[k-1];
    return tc;
}

TestCase gen_multiplicative(int n, long long k) {
    // a[i][j] = i * j
    TestCase tc;
    tc.n = n; tc.k = k;
    tc.A.assign(n+1, vector<long long>(n+1, 0));
    vector<long long> all;
    for (int i = 1; i <= n; i++)
        for (int j = 1; j <= n; j++) {
            tc.A[i][j] = (long long)i * j;
            all.push_back((long long)i * j);
        }
    sort(all.begin(), all.end());
    tc.answer = all[k-1];
    return tc;
}

TestCase gen_random_sorted(int n, long long k) {
    // Generate random sorted matrix
    TestCase tc;
    tc.n = n; tc.k = k;
    tc.A.assign(n+1, vector<long long>(n+1, 0));
    // Start with a[i][j] = i + j, then add random increments
    for (int i = 1; i <= n; i++)
        for (int j = 1; j <= n; j++)
            tc.A[i][j] = (long long)i * 1000 + (long long)j * 1000;
    // Add random noise while maintaining sorted property
    // Simple: a[i][j] = i*1000 + j*1000 + random_prefix_sum
    // Actually let's just use i*n+j with some perturbation
    // Easiest: a[i][j] = random but sorted. Generate row by row.
    // Use: a[i][j] = base[i] + cumsum of random in row i, then adjust columns
    // Simplest correct approach: a[i][j] = i*C + j*D + small_random
    // where we ensure monotonicity
    long long C = 1000000, D = 1000;
    for (int i = 1; i <= n; i++)
        for (int j = 1; j <= n; j++)
            tc.A[i][j] = (long long)i * C + (long long)j * D + (rng() % 500);
    // Fix monotonicity: ensure row-sorted and col-sorted
    for (int i = 1; i <= n; i++)
        for (int j = 2; j <= n; j++)
            tc.A[i][j] = max(tc.A[i][j], tc.A[i][j-1]);
    for (int j = 1; j <= n; j++)
        for (int i = 2; i <= n; i++)
            tc.A[i][j] = max(tc.A[i][j], tc.A[i-1][j]);

    vector<long long> all;
    for (int i = 1; i <= n; i++)
        for (int j = 1; j <= n; j++)
            all.push_back(tc.A[i][j]);
    sort(all.begin(), all.end());
    tc.answer = all[k-1];
    return tc;
}

TestCase gen_uniform(int n, long long k) {
    // a[i][j] = i + j (lots of duplicates for small n)
    // Use: a[i][j] = (i-1)*n + j for no duplicates, perfectly spread
    TestCase tc;
    tc.n = n; tc.k = k;
    tc.A.assign(n+1, vector<long long>(n+1, 0));
    // a[i][j] = i + j gives duplicates. Let's use that.
    vector<long long> all;
    for (int i = 1; i <= n; i++)
        for (int j = 1; j <= n; j++) {
            tc.A[i][j] = i + j;
            all.push_back(i + j);
        }
    sort(all.begin(), all.end());
    tc.answer = all[k-1];
    return tc;
}

TestCase gen_shifted(int n, long long k) {
    // a[i][j] = (i+n)*(j+n) - similar to interactor type 3
    TestCase tc;
    tc.n = n; tc.k = k;
    tc.A.assign(n+1, vector<long long>(n+1, 0));
    vector<long long> all;
    for (int i = 1; i <= n; i++)
        for (int j = 1; j <= n; j++) {
            tc.A[i][j] = (long long)(i + n) * (j + n);
            all.push_back(tc.A[i][j]);
        }
    sort(all.begin(), all.end());
    tc.answer = all[k-1];
    return tc;
}

// ---- Solution logic (extracted, no I/O) ----
struct Solver {
    const TestCase& tc;
    int query_count;
    vector<long long> memo;

    Solver(const TestCase& t) : tc(t), query_count(0) {
        memo.assign(2002 * 2002, -1);
    }

    long long do_query(int r, int c) {
        int key = r * 2001 + c;
        if (memo[key] != -1) return memo[key];
        query_count++;
        memo[key] = tc.A[r][c];
        return memo[key];
    }

    long long solve() {
        int n = tc.n;
        long long k = tc.k;
        long long total = (long long)n * n;

        if (n == 1) return do_query(1, 1);

        // Heap for extreme k
        long long heap_k = min(k, total - k + 1);
        if (heap_k + n <= 24000) {
            if (k <= total - k + 1) {
                priority_queue<tuple<long long, int, int>, vector<tuple<long long, int, int>>, greater<>> pq;
                vector<vector<bool>> vis(n + 1, vector<bool>(n + 1, false));
                pq.emplace(do_query(1, 1), 1, 1);
                vis[1][1] = true;
                long long result = -1;
                for (long long i = 0; i < k; i++) {
                    auto [v, r, c] = pq.top(); pq.pop();
                    result = v;
                    if (r + 1 <= n && !vis[r + 1][c]) { vis[r + 1][c] = true; pq.emplace(do_query(r + 1, c), r + 1, c); }
                    if (c + 1 <= n && !vis[r][c + 1]) { vis[r][c + 1] = true; pq.emplace(do_query(r, c + 1), r, c + 1); }
                }
                return result;
            } else {
                long long kk = total - k + 1;
                priority_queue<tuple<long long, int, int>> pq;
                vector<vector<bool>> vis(n + 1, vector<bool>(n + 1, false));
                pq.emplace(do_query(n, n), n, n);
                vis[n][n] = true;
                long long result = -1;
                for (long long i = 0; i < kk; i++) {
                    auto [v, r, c] = pq.top(); pq.pop();
                    result = v;
                    if (r - 1 >= 1 && !vis[r - 1][c]) { vis[r - 1][c] = true; pq.emplace(do_query(r - 1, c), r - 1, c); }
                    if (c - 1 >= 1 && !vis[r][c - 1]) { vis[r][c - 1] = true; pq.emplace(do_query(r, c - 1), r, c - 1); }
                }
                return result;
            }
        }

        vector<int> L(n + 1, 1), R(n + 1, n);
        long long k_rem = k;

        for (int iter = 0; iter < 100; iter++) {
            vector<int> active;
            long long total_cand = 0;
            for (int i = 1; i <= n; i++) {
                if (L[i] <= R[i]) {
                    active.push_back(i);
                    total_cand += R[i] - L[i] + 1;
                }
            }
            int na = active.size();
            if (total_cand == 0) break;
            if (total_cand == 1) {
                for (int i : active) return do_query(i, L[i]);
                break;
            }

            long long budget = 49500 - query_count;
            if (k_rem + na <= budget) {
                priority_queue<tuple<long long, int, int>, vector<tuple<long long, int, int>>, greater<>> pq;
                for (int i : active) pq.emplace(do_query(i, L[i]), i, L[i]);
                for (long long t = 1; t < k_rem; t++) {
                    auto [v, r, c] = pq.top(); pq.pop();
                    if (c + 1 <= R[r]) pq.emplace(do_query(r, c + 1), r, c + 1);
                }
                return get<0>(pq.top());
            }
            long long rev_k = total_cand - k_rem + 1;
            if (rev_k + na <= budget) {
                priority_queue<tuple<long long, int, int>> pq;
                for (int i : active) pq.emplace(do_query(i, R[i]), i, R[i]);
                for (long long t = 1; t < rev_k; t++) {
                    auto [v, r, c] = pq.top(); pq.pop();
                    if (c - 1 >= L[r]) pq.emplace(do_query(r, c - 1), r, c - 1);
                }
                return get<0>(pq.top());
            }

            // Pivot selection
            vector<long long> pvals;
            double target_frac = (double)(k_rem - 0.5) / total_cand;
            int sample_n = max(1, min(na, (int)ceil(sqrt((double)na) * 4)));
            int step = max(1, na / sample_n);
            for (int idx = 0; idx < na; idx += step) {
                int i = active[idx];
                int width = R[i] - L[i] + 1;
                int col = L[i] + (int)(target_frac * width);
                col = max(L[i], min(R[i], col));
                pvals.push_back(do_query(i, col));
            }
            sort(pvals.begin(), pvals.end());
            long long pivot = pvals[pvals.size() / 2];

            vector<int> p_le(n + 1, 0);
            {
                int j = 0;
                for (int idx = na - 1; idx >= 0; idx--) {
                    int i = active[idx];
                    j = max(j, L[i]);
                    while (j <= R[i] && do_query(i, j) <= pivot) j++;
                    p_le[i] = j - 1;
                }
            }

            long long cle = 0;
            for (int i : active) {
                int rl = min(p_le[i], R[i]);
                if (rl >= L[i]) cle += rl - L[i] + 1;
            }

            if (cle >= k_rem) {
                for (int i : active) R[i] = min(R[i], p_le[i]);
            } else {
                k_rem -= cle;
                for (int i : active) L[i] = max(L[i], p_le[i] + 1);
            }
        }
        return -1; // shouldn't reach
    }
};

int main() {
    struct TestDef {
        string name;
        function<TestCase()> gen;
    };

    vector<TestDef> tests;

    // Small tests
    tests.push_back({"additive n=100 k=5000", []{ return gen_additive(100, 5000); }});
    tests.push_back({"multiplicative n=100 k=5000", []{ return gen_multiplicative(100, 5000); }});

    // Medium tests
    tests.push_back({"additive n=500 k=125000", []{ return gen_additive(500, 125000); }});
    tests.push_back({"multiplicative n=500 k=125000", []{ return gen_multiplicative(500, 125000); }});
    tests.push_back({"random n=500 k=125000", []{ return gen_random_sorted(500, 125000); }});

    // Hard tests (n=2000)
    tests.push_back({"additive n=2000 k=2000000", []{ return gen_additive(2000, 2000000); }});
    tests.push_back({"multiplicative n=2000 k=2000000", []{ return gen_multiplicative(2000, 2000000); }});
    tests.push_back({"shifted n=2000 k=2000000", []{ return gen_shifted(2000, 2000000); }});
    tests.push_back({"random n=2000 k=2000000", []{ return gen_random_sorted(2000, 2000000); }});

    // Extreme k
    tests.push_back({"multiplicative n=2000 k=1", []{ return gen_multiplicative(2000, 1); }});
    tests.push_back({"multiplicative n=2000 k=4000000", []{ return gen_multiplicative(2000, 4000000); }});
    tests.push_back({"multiplicative n=2000 k=100000", []{ return gen_multiplicative(2000, 100000); }});
    tests.push_back({"multiplicative n=2000 k=3900000", []{ return gen_multiplicative(2000, 3900000); }});

    for (auto& t : tests) {
        auto tc = t.gen();
        Solver s(tc);
        long long result = s.solve();
        bool correct = (result == tc.answer);
        double score;
        int n = tc.n;
        int used = s.query_count;
        if (!correct) score = 0.0;
        else if (used <= n) score = 1.0;
        else if (used >= 50000) score = 0.0;
        else score = (50000.0 - used) / (50000.0 - n);

        printf("%-45s n=%4d k=%8lld queries=%6d correct=%s score=%.4f\n",
               t.name.c_str(), tc.n, tc.k, used, correct ? "YES" : "NO", score);
    }

    return 0;
}