#include "tickets.h"
#include <vector>
#include <cmath>
#include <algorithm>
#include <queue>
using namespace std;
#define F first
#define S second
#define PB push_back
using ll = long long;
// each vector<int> in x is sorted
long long find_maximum(int k, vector<vector<int>> x) {
/* sample interaction
int n = x.size();
int m = x[0].size();
vector<vector<int>> answer;
for (int i = 0; i < n; i++) {
vector<int> row(m);
for (int j = 0; j < m; j++) {
if (j < k) {
row[j] = j;
} else {
row[j] = -1;
}
}
answer.push_back(row);
}
allocate_tickets(answer);
return 1;
*/
int n = x.size();
int m = x[0].size();
vector<vector<int>> answer(n, vector<int>(m, -1));
if(m == 1) {
// subtask 1 (only a single way to pick everything)
vector<int> A;
for(int i = 0; i < n; i++) {
A.PB(x[i][0]);
answer[i][0] = 0;
}
sort(A.begin(), A.end());
ll sum = 0;
for(int a : A) {
sum += abs(a-A[n/2]);
}
allocate_tickets(answer);
return sum;
}
else {
// same sol as subtask 2 but repeat k times
ll sum = 0;
struct Info {
int val;
int i, j, k;
bool operator<(const Info &rhs) const {
return val < rhs.val;
}
};
for(int q = 0; q < k; q++) {
priority_queue<Info> pq;
for(int i = 0; i < n; i++) {
int l = -1, r = -1;
for(int j = 0; j < m; j++) {
if(answer[i][j] != -1) continue;
if(l == -1) l = j;
r = j;
}
answer[i][l] = q;
sum -= x[i][l];
pq.push({x[i][l]+x[i][r], i, l, r});
}
for(int i = 0; i < n/2; i++) {
answer[pq.top().i][pq.top().j] = -1;
answer[pq.top().i][pq.top().k] = q;
sum += pq.top().val;
pq.pop();
}
}
allocate_tickets(answer);
return sum;
}
allocate_tickets(answer);
return -1;
}
