File size: 9,381 Bytes
1fd0050
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
#include <bits/stdc++.h>
using namespace std;

struct TestCase {
    int n;
    long long k;
    vector<vector<long long>> A;
    long long answer;
};

mt19937_64 rng_gen(42);

TestCase gen_matrix(int n, long long k, function<long long(int,int)> valfn) {
    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] = valfn(i, j);
            all.push_back(tc.A[i][j]);
        }
    sort(all.begin(), all.end());
    tc.answer = all[k-1];
    return tc;
}

TestCase gen_multiplicative(int n, long long k) {
    return gen_matrix(n, k, [](int i, int j) -> long long { return (long long)i * j; });
}
TestCase gen_shifted(int n, long long k) {
    return gen_matrix(n, k, [n](int i, int j) -> long long { return (long long)(i + n) * (j + n); });
}
TestCase gen_additive(int n, long long k) {
    return gen_matrix(n, k, [](int i, int j) -> long long { return i + j; });
}
TestCase gen_random_sorted(int n, long long k) {
    TestCase tc;
    tc.n = n; tc.k = k;
    tc.A.assign(n+1, vector<long long>(n+1, 0));
    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_gen() % 500);
    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;
}

struct Solver {
    const TestCase& tc;
    int query_count;
    vector<long long> memo;
    int n;

    Solver(const TestCase& t) : tc(t), query_count(0), n(t.n) {
        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() {
        long long k = tc.k;
        long long N2 = (long long)n * n;

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

        long long heap_k = min(k, N2 - k + 1);
        if (heap_k + n <= 24000) {
            if (k <= N2 - k + 1) {
                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);
                }
                return 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);
                }
                return result;
            }
        }

        // Quickselect with improved pivot selection
        vector<int> L(n + 1, 1), R(n + 1, n);
        long long k_rem = k;
        mt19937_64 rng(12345);

        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());
            }

            // Improved pivot selection: uniform random sampling from candidate area
            // Build prefix sums of widths
            vector<long long> pref(na + 1, 0);
            for (int idx = 0; idx < na; idx++) {
                int i = active[idx];
                pref[idx + 1] = pref[idx] + (R[i] - L[i] + 1);
            }

            int num_samples = min(15, na);
            vector<long long> samp;
            for (int s = 0; s < num_samples; s++) {
                long long pick = 1 + rng() % total_cand;
                int idx = (int)(lower_bound(pref.begin() + 1, pref.end(), pick) - pref.begin()) - 1;
                idx = max(0, min(na - 1, idx));
                int i = active[idx];
                int width = R[i] - L[i] + 1;
                int col = L[i] + rng() % width;
                samp.push_back(do_query(i, col));
            }
            sort(samp.begin(), samp.end());

            // Pick at quantile matching k_rem/total_cand
            double q = (double)(k_rem - 0.5) / total_cand;
            int pidx = (int)(q * (double)(samp.size() - 1));
            pidx = max(0, min((int)samp.size() - 1, pidx));
            long long pivot = samp[pidx];

            // Staircase walk (bottom-to-top)
            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;
    }
};

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

    vector<TestDef> 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); }});
    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); }});
    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); }});
    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 used = s.query_count;
        if (!correct) score = 0.0;
        else if (used <= tc.n) score = 1.0;
        else if (used >= 50000) score = 0.0;
        else score = (50000.0 - used) / (50000.0 - tc.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;
}