#include "tickets.h"
#include <vector>
#include <algorithm>
#include <queue>
#include <cassert>
using namespace std;
typedef long long LL;
typedef vector<int> IVec;
typedef pair<LL, LL> pii;
long long find_maximum(int k, vector<IVec> d) {
int N = d.size(), M = d[0].size(); // N color ×M tickets
vector<IVec> o(N); // o.resize(N);
for (int i = 0; i < N; i++) {
for (int j = 0; j < M; j++) o[i].push_back(j); // o[N][i]:index of ticket which is ith min from left
sort(o[i].begin(), o[i].end(), [&](int a, int b) { return d[i][a] < d[i][b]; });
}
auto dval = [&](int c, int i) { return d[c][o[c][i]]; }; // num of ticket of color c which is ith form left
LL cost = 0;
IVec Plus(N, 0); // Plus[c] => How many '+'s the color c will send
priority_queue<pii> gain; // gain: (change to cost by turning a - into a +, color index)
for (int c = 0; c < N; c++) {
for (int j = 0; j < k; j++) cost -= dval(c, j); // First incur cost of all L_c
assert(Plus[c] == 0);
gain.push(make_pair(dval(c, M - 1 - Plus[c]) + dval(c, k - 1 - Plus[c]), c)); // Queue a possible - to +
}
for (int i = 0; i < N * k / 2; i++) { // Take ck/2 +s
pii it = gain.top(); gain.pop();
cost += it.first;
int c = it.second, &pc = Plus[c]; // If the color hasnt reached k +s, queue another possible +
if (++pc < k) gain.push(pii(dval(c, M - 1 - pc) + dval(c, k - 1 - pc), c));
}
while (!gain.empty()) gain.pop();
for (int c = 0; c < N; c++) gain.push(pii(Plus[c], c)); // gain: (+s remaining, color index)
IVec minus(N, 0); // minus[color] => How many -s the color has sent so far
vector<IVec> answer(N, IVec(M, -1));
IVec take(N, 0); // take[c] => Does this c send a + in this rounds
for (int i = 0; i < k; i++) { // Simulate k days of tickets
for (int c = 0; c < N / 2; c ++) take[gain.top().second] = 1, gain.pop(); // Pick N/2 colors with the most +s left
for (int c = 0; c < N; c++) { // Check if each color sent a - or +, and change answer accordingly
int &p = Plus[c], &m = minus[c]; // color c sent a + this round, send the smallest +, Queue the color back
if (take[c]) answer[c][o[c][M - p]] = i, gain.push(pii(--p, c));
else answer[c][o[c][m++]] = i; // color c sent a - this round, send the smallest -
take[c] = 0;
}
}
allocate_tickets(answer);
return cost;
}
