docker_space/frontier_cs_4/solution_new.cpp

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

int n;
long long k;
vector<long long> memo;
int query_count = 0;

long long do_query(int r, int c) {
    int key = r * 2001 + c;
    if (memo[key] != -1) return memo[key];
    cout << "QUERY " << r << " " << c << "\n";
    cout.flush();
    long long v;
    cin >> v;
    memo[key] = v;
    query_count++;
    return v;
}

void done(long long ans) {
    cout << "DONE " << ans << "\n";
    cout.flush();
    exit(0);
}

// Staircase walk: count elements <= x, track boundary per row
// posOut[i] = max column j such that a[i][j] <= x (0 if none)
// Walk top-to-bottom, right-to-left: O(n) new queries (with caching, even less)
long long count_leq(long long x, vector<int>& pos) {
    pos.assign(n + 1, 0);
    long long cnt = 0;
    int j = n;
    for (int i = 1; i <= n; i++) {
        while (j >= 1 && do_query(i, j) > x) j--;
        pos[i] = j;
        cnt += j;
        if (j == 0) {
            for (int ii = i + 1; ii <= n; ii++) pos[ii] = 0;
            break;
        }
    }
    return cnt;
}

long long count_lt(long long x, vector<int>& pos) {
    pos.assign(n + 1, 0);
    long long cnt = 0;
    int j = n;
    for (int i = 1; i <= n; i++) {
        while (j >= 1 && do_query(i, j) >= x) j--;
        pos[i] = j;
        cnt += j;
        if (j == 0) {
            for (int ii = i + 1; ii <= n; ii++) pos[ii] = 0;
            break;
        }
    }
    return cnt;
}

// Bounded staircase walk: only check within posLow[i]+1..posHigh[i]
// More query-efficient when bounds are tight
long long count_leq_bounded(long long x, vector<int>& pos, const vector<int>& lo, const vector<int>& hi) {
    pos.assign(n + 1, 0);
    long long cnt = 0;
    // For row i, we know answer is in [lo[i]+1, hi[i]]
    // Walk top to bottom. j starts from hi[1] and decreases.
    int j = n; // will be clamped
    for (int i = 1; i <= n; i++) {
        j = min(j, hi[i]);
        if (j < 1) {
            pos[i] = 0;
            for (int ii = i + 1; ii <= n; ii++) pos[ii] = 0;
            break;
        }
        while (j >= 1 && do_query(i, j) > x) j--;
        pos[i] = j;
        cnt += max(0, j);
        if (j == 0) {
            for (int ii = i + 1; ii <= n; ii++) pos[ii] = 0;
            break;
        }
    }
    return cnt;
}

void solve() {
    long long N2 = (long long)n * n;

    if (n == 1) { done(do_query(1, 1)); return; }
    if (k == 1) { done(do_query(1, 1)); return; }
    if (k == N2) { done(do_query(n, n)); return; }

    // For small k or k close to N2, use heap directly
    long long heap_k = min(k, N2 - k + 1);
    if (heap_k + n <= 24000) {
        if (k <= N2 - k + 1) {
            // k-th smallest via min-heap on rows
            priority_queue<tuple<long long, int, int>, vector<tuple<long long, int, int>>, greater<>> pq;
            for (int i = 1; i <= n; i++)
                pq.emplace(do_query(i, 1), i, 1);
            long long result = -1;
            for (long long t = 0; t < k; t++) {
                auto [v, r, c] = pq.top(); pq.pop();
                result = v;
                if (c + 1 <= n) pq.emplace(do_query(r, c + 1), r, c + 1);
            }
            done(result);
        } else {
            long long kk = N2 - k + 1;
            priority_queue<tuple<long long, int, int>> pq;
            for (int i = 1; i <= n; i++)
                pq.emplace(do_query(i, n), i, n);
            long long result = -1;
            for (long long t = 0; t < kk; t++) {
                auto [v, r, c] = pq.top(); pq.pop();
                result = v;
                if (c - 1 >= 1) pq.emplace(do_query(r, c - 1), r, c - 1);
            }
            done(result);
        }
    }

    // Value-based narrowing approach
    mt19937_64 rng(12345);

    // Initialize: posHigh from count_leq(a[n][n]), posLow all 0
    vector<int> posHigh(n + 1, n), posLow(n + 1, 0);
    long long cntHigh = N2, cntLow = 0;
    long long highTh, lowTh;

    // Get initial bounds - use anti-diagonal estimate
    // Row ceil(k/n): a[ceil(k/n)][n] is an upper bound
    int rBound = (int)min((long long)n, (k + n - 1) / n);
    highTh = do_query(rBound, n);
    cntHigh = count_leq(highTh, posHigh);
    if (cntHigh < k) {
        highTh = do_query(n, n);
        cntHigh = count_leq(highTh, posHigh);
    }

    // Lower bound: a[floor(k/n)+1][1] or just 0
    // Actually, row floor(k/n): a[floor(k/n)][1] has at most floor(k/n)*n elements <= it from first floor(k/n) rows
    // But it's more complex. Just use cntLow = 0 initially.
    lowTh = do_query(1, 1) - 1; // everything is >= a[1][1]
    // posLow stays all 0

    const int SAMPLES = 15;

    for (int iter = 0; iter < 200; iter++) {
        if (cntHigh < k) {
            highTh = do_query(n, n);
            cntHigh = count_leq(highTh, posHigh);
            if (cntHigh < k) break; // shouldn't happen
        }

        long long candTotal = cntHigh - cntLow;
        if (candTotal <= 0) { done(highTh); }

        long long needSmall = k - cntLow;
        long long needLarge = cntHigh - k + 1;

        // Count non-empty rows and build prefix sums for sampling
        vector<long long> pref(n + 1, 0);
        int nonempty = 0;
        for (int i = 1; i <= n; i++) {
            int len = posHigh[i] - posLow[i];
            if (len > 0) nonempty++;
            pref[i] = pref[i - 1] + max(0, len);
        }
        candTotal = pref[n];
        if (candTotal <= 0) { done(highTh); }

        // Check if we can enumerate directly
        long long budget = 49500 - query_count;
        long long enumCost = min(needSmall, needLarge) + nonempty + 10;
        if (enumCost <= budget) {
            // Enumerate via heap
            if (needSmall <= needLarge) {
                priority_queue<tuple<long long, int, int>, vector<tuple<long long, int, int>>, greater<>> pq;
                for (int i = 1; i <= n; i++) {
                    int L = posLow[i] + 1;
                    int R = posHigh[i];
                    if (L >= 1 && L <= n && L <= R) {
                        pq.emplace(do_query(i, L), i, L);
                    }
                }
                long long result = 0;
                for (long long t = 0; t < needSmall; t++) {
                    auto [v, r, c] = pq.top(); pq.pop();
                    result = v;
                    if (c + 1 <= posHigh[r]) pq.emplace(do_query(r, c + 1), r, c + 1);
                }
                done(result);
            } else {
                priority_queue<tuple<long long, int, int>> pq;
                for (int i = 1; i <= n; i++) {
                    int L = posLow[i] + 1;
                    int R = posHigh[i];
                    if (R >= 1 && R <= n && L <= R) {
                        pq.emplace(do_query(i, R), i, R);
                    }
                }
                long long result = 0;
                for (long long t = 0; t < needLarge; t++) {
                    auto [v, r, c] = pq.top(); pq.pop();
                    result = v;
                    int L = posLow[r] + 1;
                    if (c - 1 >= L) pq.emplace(do_query(r, c - 1), r, c - 1);
                }
                done(result);
            }
        }

        // Budget check for another iteration
        if (budget < 2 * n + SAMPLES + 200) {
            // Try to enumerate anyway
            if (nonempty + min(needSmall, needLarge) + 10 <= budget) {
                // same enumeration as above
                if (needSmall <= needLarge) {
                    priority_queue<tuple<long long, int, int>, vector<tuple<long long, int, int>>, greater<>> pq;
                    for (int i = 1; i <= n; i++) {
                        int L = posLow[i] + 1, R = posHigh[i];
                        if (L >= 1 && L <= n && L <= R) pq.emplace(do_query(i, L), i, L);
                    }
                    long long result = 0;
                    for (long long t = 0; t < needSmall; t++) {
                        auto [v, r, c] = pq.top(); pq.pop();
                        result = v;
                        if (c + 1 <= posHigh[r]) pq.emplace(do_query(r, c + 1), r, c + 1);
                    }
                    done(result);
                } else {
                    priority_queue<tuple<long long, int, int>> pq;
                    for (int i = 1; i <= n; i++) {
                        int L = posLow[i] + 1, R = posHigh[i];
                        if (R >= 1 && R <= n && L <= R) pq.emplace(do_query(i, R), i, R);
                    }
                    long long result = 0;
                    for (long long t = 0; t < needLarge; t++) {
                        auto [v, r, c] = pq.top(); pq.pop();
                        result = v;
                        if (c - 1 >= posLow[r] + 1) pq.emplace(do_query(r, c - 1), r, c - 1);
                    }
                    done(result);
                }
            }
            done(highTh); // fallback
        }

        // Sample pivot values from candidate segments
        vector<long long> samp;
        samp.reserve(SAMPLES);
        for (int s = 0; s < SAMPLES; s++) {
            long long pick = 1 + rng() % candTotal;
            int row = (int)(lower_bound(pref.begin() + 1, pref.end(), pick) - pref.begin());
            if (row < 1 || row > n) continue;
            int len = posHigh[row] - posLow[row];
            if (len <= 0) { continue; }
            int col = posLow[row] + 1 + rng() % len;
            col = max(1, min(n, col));
            samp.push_back(do_query(row, col));
        }
        if (samp.empty()) { done(highTh); }

        sort(samp.begin(), samp.end());

        // Pick pivot at the quantile matching needSmall/candTotal
        double q = (double)needSmall / (double)candTotal;
        int idx = (int)(q * (double)(samp.size() - 1));
        idx = max(0, min((int)samp.size() - 1, idx));
        long long pivot = samp[idx];

        // Count elements <= pivot
        vector<int> posPivot;
        long long cntPivot = count_leq(pivot, posPivot);

        if (cntPivot >= k) {
            // Tighten high
            if (pivot == highTh && cntPivot == cntHigh) {
                // Pivot equals current high and same count - duplicates
                vector<int> posLess;
                long long cntLess = count_lt(pivot, posLess);
                if (cntLess < k) {
                    done(pivot); // k-th element equals pivot
                }
                // Narrow to strictly less than pivot
                highTh = pivot - 1;
                posHigh = posLess;
                cntHigh = cntLess;
            } else {
                highTh = pivot;
                posHigh = posPivot;
                cntHigh = cntPivot;
            }
        } else {
            // Tighten low
            lowTh = pivot;
            posLow = posPivot;
            cntLow = cntPivot;
        }
    }
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cin >> n >> k;
    memo.assign(2002 * 2002, -1);
    solve();
    return 0;
}