#include "tickets.h"
#include <vector>
#include <algorithm>
#include <queue>
using namespace std;
/**
* IOI 2020 - Carnival Tickets
* To'liq yechim: O(NM log N)
*/
long long find_maximum(int k, vector<vector<int>> x) {
int n = x.size();
int m = x[0].size();
vector<int> L(n, k), R(n, m - 1);
priority_queue<pair<long long, int>> pq;
long long total_prize = 0;
// 1. Dastlabki holat: har bir rangdan eng kichik k ta chiptani olamiz
for (int i = 0; i < n; i++) {
for (int j = 0; j < k; j++) {
total_prize -= x[i][j];
}
// Kichikni kattaga almashtirishdagi foyda
pq.push({(long long)x[i][m - 1] + x[i][k - 1], i});
}
// 2. n*k/2 marta eng foydali almashtirishlarni bajaramiz
for (int i = 0; i < (n * k) / 2; i++) {
pair<long long, int> top = pq.top();
pq.pop();
total_prize += top.first;
int id = top.second;
L[id]--;
R[id]--;
if (L[id] > 0) {
pq.push({(long long)x[id][R[id]] + x[id][L[id] - 1], id});
}
}
// 3. Chiptalarni raundlarga taqsimlash
vector<vector<int>> s(n, vector<int>(m, -1));
vector<int> rounds_count(k, 0);
vector<pair<int, int>> colors;
for (int i = 0; i < n; i++) {
colors.push_back({L[i], i});
}
for (int r = 0; r < k; r++) {
sort(colors.rbegin(), colors.rend());
for (int i = 0; i < n; i++) {
int id = colors[i].second;
if (i < n / 2 && L[id] > 0) {
// Kichik chiptani ishlatish
s[id][k - L[id]] = r; // x[id][0...k-1] oralig'idagi chipta
L[id]--;
colors[i].first--;
} else {
// Katta chiptani ishlatish
s[id][m - 1 - (k - 1 - L[id] - rounds_count[r])] = r;
// Bu qismda m-1 dan boshlab chiptalar olinadi
}
}
}
// Taqsimotni qayta ishlash (soddalashtirilgan mantiq)
vector<vector<int>> final_s(n, vector<int>(m, -1));
vector<int> cur_l(n, 0), cur_r(n, m - 1);
vector<int> used_as_plus(n, k - L[n-1]); // Nechta katta chipta olingani
// To'g'ri taqsimlash algoritmi
vector<int> pos_count(n);
for(int i=0; i<n; i++) pos_count[i] = k - L[i];
for (int r = 0; r < k; r++) {
vector<int> take_pos;
vector<pair<int, int>> temp;
for(int i=0; i<n; i++) temp.push_back({pos_count[i], i});
sort(temp.rbegin(), temp.rend());
for(int i=0; i<n; i++) {
int id = temp[i].second;
if (i < n/2) {
final_s[id][cur_r[id]--] = r;
pos_count[id]--;
} else {
final_s[id][cur_l[id]++] = r;
}
}
}
allocate_tickets(final_s);
return total_prize;
}
