File size: 9,083 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
#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));
    for (int i = 1; i <= n; i++) for (int j = 1; j <= n; j++) tc.A[i][j] = (long long)i * 1000000 + (long long)j * 1000 + (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;
    // diagnostics
    vector<double> split_ratios;

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

        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;
                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 = N2 - 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: sample from ALL active rows at target_frac position
            // Use WEIGHTED quantile selection
            double target_frac = (double)(k_rem - 0.5) / total_cand;

            // Sample from ALL active rows
            vector<pair<long long, int>> vw; // (value, weight=width)
            for (int i : active) {
                int width = R[i] - L[i] + 1;
                int col = L[i] + (int)(target_frac * (width - 1));
                col = max(L[i], min(R[i], col));
                vw.push_back({do_query(i, col), width});
            }

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

            // WEIGHTED quantile: find value at target_frac of total weight
            long long total_w = 0;
            for (auto& [v, w] : vw) total_w += w;
            long long target_w = max(1LL, (long long)(target_frac * total_w));
            long long cum = 0;
            long long pivot = vw.back().first;
            for (auto& [v, w] : vw) {
                cum += w;
                if (cum >= target_w) { pivot = v; break; }
            }

            // 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;
            }

            double ratio = (double)min(cle, total_cand - cle) / total_cand;
            split_ratios.push_back(ratio);

            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({"mult 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({"mult 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({"mult 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({"mult n=2000 k=100000", []{ return gen_multiplicative(2000, 100000); }});
    tests.push_back({"mult 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);
        int used = s.query_count;
        double score = !correct ? 0.0 : (used <= tc.n ? 1.0 : (used >= 50000 ? 0.0 : (50000.0 - used) / (50000.0 - tc.n)));
        printf("%-45s q=%6d %s score=%.4f  splits:", t.name.c_str(), used, correct ? "OK" : "WRONG", score);
        for (double r : s.split_ratios) printf(" %.3f", r);
        printf("\n");
    }
}