#include "tickets.h"
#include <bits/stdc++.h>
using namespace std;
#define int long long
#define sig signed
#define arr array
#define vec vector
#define pii pair<int, int>
#define fir first
#define sec second
#define mp make_pair
const int N = 1505, M = 1505, INF = 1e18;
int n, m, k;
arr<arr<int, M>, N> a; // sorted in input
arr<int, N> blc;
void blc_cmp() {
for (int i = 1; i <= n; i++)
blc[i] = (i <= n / 2) ? k : 0;
while (true) {
vec<int> bst = {0, INF, INF};
for (int i = 1; i <= n; i++) {
for (int j = 1; j <= n; j++) {
if (i == j) continue;
if (blc[i] == 0 || blc[j] == k) continue;
int x = a[i][blc[i]] + a[j][m - k + blc[i]];
int y = -a[j][blc[j] + 1] - a[j][m - k + blc[j] + 1];
bst = max(bst, {x + y, i, j});
}
}
if (bst[0] == 0) break;
blc[bst[1]]--, blc[bst[2]]++;
}
// for (int i = 1; i <= n; i++)
// cout << i << ": " << blc[i] << '\n';
}
arr<arr<int, M>, N> ass;
int ass_cmp() {
int ans = 0;
for (int x = 1; x <= k; x++) {
int y = k - x + 1;
vec<pii> ord;
for (int i = 1; i <= n; i++) {
if (m - y + blc[i] + 1 > m) ord.push_back({-INF, i});
else ord.push_back({a[i][m - y + blc[i] + 1], i});
}
sort(ord.rbegin(), ord.rend());
for (int j = 0; j < n; j++) {
int i = ord[j].sec;
if (j <= n / 2 - 1) {
ass[i][m - y + blc[i] + 1] = x;
ans += a[i][m - y + blc[i] + 1];
} else {
ass[i][blc[i]] = x;
ans -= a[i][blc[i]];
blc[i]--;
}
}
}
return ans;
}
// Return indices 0-indexed
int find_maximum(sig _k, vec<vec<sig>> _a) {
n = _a.size(), m = _a[0].size(), k = _k;
for (int i = 1; i <= n; i++)
for (int j = 1; j <= m; j++)
a[i][j] = _a[i - 1][j - 1];
blc_cmp();
int chck = 0;
for (int i = 1; i <= n; i++)
for (int j = 1; j <= m; j++) {
if (j <= blc[i]) chck -= a[i][j];
if (j >= m - k + blc[i] + 1) chck += a[i][j];
}
int ans = ass_cmp();
assert(ans == chck);
vec<vec<sig>> x;
for (int i = 1; i <= n; i++) {
vec<sig> y;
for (int j = 1; j <= m; j++) y.push_back(ass[i][j] - 1);
x.push_back(y);
}
allocate_tickets(x);
return ans;
}
