#include "tickets.h"
#include <bits/stdc++.h>
using namespace std;
vector<vector<int>> ans;
// 1 5
// 10 10
// 8 12
// 12 20
// a1 a2 b1 c1
// b2 c2 d1 d2
// -a1 -b1 -c1 d2
// -a1 -b1 +c2 +d2
vector<vector<int>> p;
vector<vector<bool>> used;
long long bst;
void rec(int k, int at, vector<vector<int>> &x) {
if (at == p.size()) {
vector<int> v;
for (int i = 0; i < p.size(); i++) {
for (auto idx: p[i]) {
v.push_back(idx);
}
}
sort(v.begin(), v.end());
long long cur = 0;
for (int i = 0; i < v.size() / 2; i++)
cur += v[v.size() - 1 - i] - v[i];
bst = max(bst, cur);
return;
}
if (p[at].size() == k) {
rec(k, at + 1, x);
return;
}
for (int i = 0; i < x[at].size(); i++)
if (!used[at][i]) {
used[at][i] = true;
p[at].push_back(x[at][i]);
rec(k, at, x);
p[at].pop_back();
used[at][i] = false;
}
}
long long slow(int k, vector<vector<int>> x) {
int n = x.size();
int m = x[0].size();
p = vector<vector<int>>(n);
used = vector<vector<bool>>(n, vector<bool>(m, false));
bst = -1;
rec(k, 0, x);
return bst;
}
long long fast(int k, vector<vector<int>> x) {
int n = x.size();
int m = x[0].size();
ans = vector<vector<int>>(n, vector<int>(m, -1));
long long res = 0;
vector<pair<int, pair<int, int>>> all;
bool binary = true;
for (int i = 0; i < n; i++)
for (int j = 0; j < m; j++) {
all.push_back({x[i][j], {i, j}});
if (x[i][j] > 1) binary = false;
}
sort(all.begin(), all.end());
vector<vector<pair<int, int>>> high(n), pickedLow(n), pickedHigh(n);
int median = (all[n * m / 2].first + all[n * m / 2 - 1].first) / 2;
int cnt = 0;
for (int i = 0; i < all.size(); i++) {
int color = all[i].second.first;
int ticket = all[i].second.second;
if (i < all.size() / 2) {
if (pickedLow[color].size() < k) {
pickedLow[color].push_back({median - all[i].first, ticket});
cnt++;
res += median - all[i].first;
}
} else {
high[color].push_back({all[i].first - median, ticket});
}
}
for (int i = 0; i < n; i++) {
while (pickedLow[i].size() + pickedHigh[i].size() < k) {
res += high[i].back().first;
pickedHigh[i].push_back(high[i].back());
high[i].pop_back();
}
}
vector<vector<long long>> b(n);
int need = cnt - n * k / 2;
for (int i = 0; i < n; i++) {
long long s = 0;
b[i].push_back(0);
for (int j = 1; j <= min(pickedLow[i].size(), high[i].size()); j++) {
s += high[i][high[i].size() - j].first - pickedLow[i][pickedLow[i].size() - j].first;
b[i].push_back(s);
}
}
long long inf = -(1LL << 62);
vector<vector<long long>> dp(n + 1, vector<long long>(need + 1, inf));
dp[0][0] = 0;
vector<vector<int>> back(n + 1, vector<int>(need + 1));
for (int i = 0; i < n; i++) {
for (int j = 0; j <= need; j++) {
if (dp[i][j] == inf) continue;
for (int k = 0; k < b[i].size(); k++) {
if (j + k > need) break;
if (dp[i][j] + b[i][k] > dp[i + 1][k + j]) {
dp[i + 1][k + j] = dp[i][j] + b[i][k];
back[i + 1][k + j] = k;
}
}
}
}
res += dp[n][need];
for (int i = n - 1; i >= 0; i--) {
int k = back[i + 1][need];
need -= k;
for (int j = 0; j < k; j++) {
pickedLow[i].pop_back();
pickedHigh[i].push_back(high[i].back());
high[i].pop_back();
}
}
// vector<int> at(n, 0);
// while (need--) {
// int idx = -1, best = 0;
// for (int i = 0; i < n; i++) {
// if (at[i] == b[i].size()) continue;
// int cur;
// cur = b[i][at[i]];
// if (idx == -1 || cur > best) {
// best = cur;
// idx = i;
// }
// }
// assert(idx != -1);
// // switch
// res += best;
// at[idx]++;
// pickedLow[idx].pop_back();
// pickedHigh[idx].push_back(high[idx].back());
// high[idx].pop_back();
// }
// construct
for (int r = 0; r < k; r++) {
vector<bool> used(n);
vector<pair<int, int>> counts;
for (int i = 0; i < n; i++) counts.push_back({pickedLow[i].size(), i});
sort(counts.rbegin(), counts.rend());
int taken = 0;
for (int ii = 0; ii < n; ii++) {
int i = counts[ii].second;
if (pickedLow[i].size() > 0) {
used[i] = true;
taken++;
ans[i][pickedLow[i].back().second] = r;
pickedLow[i].pop_back();
}
if (taken == n / 2) break;
}
taken = 0;
for (int i = 0; i < n; i++) {
if (pickedHigh[i].size() > 0 && !used[i]) {
taken++;
ans[i][pickedHigh[i].back().second] = r;
pickedHigh[i].pop_back();
}
if (taken == n / 2) break;
}
}
allocate_tickets(ans);
long long sl = slow(k, x);
assert(sl >= res);
return sl;
}
long long find_maximum(int k, std::vector<std::vector<int>> x) {
return fast(k, x);
}
// a b a c
// c b d d
