import java.util.Arrays;
import java.util.PriorityQueue;
public class tickets {
long find_maximum(int k, int[][] x) {
int n = x.length;
int m = x[0].length;
int[][] answer = new int[n][m];
for (int i = 0; i < n; i++) Arrays.fill(answer[i], -1);
long ans = 0;
// i -> m - k + i
PriorityQueue<int[]> pq = new PriorityQueue<int[]>(
(int[] a, int[] b) -> {
int col1 = a[0], i1 = a[1];
int col2 = b[0], i2 = b[1];
int s1 = x[col1][i1] + x[col1][m - k + i1];
int s2 = x[col2][i2] + x[col2][m - k + i2];
return s1 != s2 ? s1 - s2 : col1 - col2;
});
for (int i = 0; i < n; i++) {
pq.add(new int[] { i, 0 });
}
int[] until = new int[n];
for (int s = 0; s < k * (n / 2); s++) {
int[] p = pq.poll();
int col = p[0];
int ind = p[1];
until[col]++;
if (ind < k - 1) {
p[1]++;
pq.add(p);
}
}
Tickets[] tickets = new Tickets[n];
for (int i = 0; i < n; i++) {
tickets[i] = new Tickets(i, x[i], until[i] - 1, m - k + until[i]);
}
for (int i = 0; i < k; i++) {
Arrays.sort(tickets, (Tickets a, Tickets b) -> Math.max(b.numMinus(), b.numPlus()) -
Math.max(a.numMinus(), a.numPlus()));
int balance = 0;
int numDaysLeft = k - i;
for (int j = 0; j < n; j++) {
Tickets t = tickets[j];
int col = t.color;
if (t.numMinus() == numDaysLeft) {
int index = t.minusTill;
t.minusTill--;
answer[col][index] = i;
balance--;
} else if (t.numPlus() == numDaysLeft) {
int index = t.plusAfter;
t.plusAfter++;
answer[col][index] = i;
balance++;
} else {
if (t.numMinus() == 0) {
int index = t.plusAfter;
t.plusAfter++;
answer[col][index] = i;
balance++;
} else if (t.numPlus() == 0) {
int index = t.minusTill;
t.minusTill--;
answer[col][index] = i;
balance--;
} else {
if (balance <= 0) {
int index = t.plusAfter;
t.plusAfter++;
answer[col][index] = i;
balance++;
} else {
int index = t.minusTill;
t.minusTill--;
answer[col][index] = i;
balance--;
}
}
}
}
}
for (int i = 0; i < n; i++) {
for (int j = 0; j < until[i]; j++) {
ans -= x[i][j];
}
for (int j = until[i] + m - k; j < m; j++) {
ans += x[i][j];
}
}
grader.allocate_tickets(answer);
return ans;
}
long doTheJob(int[][] x, int k, int n, int m) {
Tickets[] tickets = new Tickets[n];
int[][] answer = new int[n][m];
int[] until = new int[n], after = new int[n];
// triples [color, index]
PriorityQueue<int[]> pq = new PriorityQueue<int[]>(
(int[] a, int[] b) -> {
int col1 = a[0], i1 = a[1];
int col2 = b[0], i2 = b[1];
return x[col1][i1] != x[col2][i2] ? Integer.compare(x[col1][i1], x[col2][i2]) : col1 - col2;
}
);
for (int i = 0; i < n; i++) {
pq.add(new int[] { i, 0 });
}
for (int i = 0; i < k * (n / 2); i++) {
int[] p = pq.poll();
int col = p[0];
until[col]++;
p[1]++;
if (p[1] < k) {
pq.add(p);
}
}
pq = new PriorityQueue<>(
(int[] a, int[] b) -> {
int col1 = a[0], i1 = a[1];
int col2 = b[0], i2 = b[1];
return x[col1][i1] != x[col2][i2] ? -Integer.compare(x[col1][i1], x[col2][i2]) : -col1 + col2;
}
);
Arrays.fill(after, m);
for (int i = 0; i < n; i++) {
if (until[i] >= m) {
continue;
}
pq.add(new int[] { i, m - 1 });
}
for (int i = 0; i < k * (n / 2); i++) {
int[] p = pq.poll();
int col = p[0];
if (m - p[1] + until[col] > k) continue;
after[col] = p[1];
p[1]--;
if (p[1] >= until[col]) {
pq.add(p);
}
}
for (int i = 0; i < n; i++) {
tickets[i] = new Tickets(i, x[i], until[i] - 1, after[i]);
//System.out.println("Ticket " + i + ": " + (until[i] - 1) + " " + after[i]);
}
for (int i = 0; i < k; i++) {
Arrays.sort(tickets, (Tickets a, Tickets b) -> Math.max(b.numMinus(), b.numPlus()) -
Math.max(a.numMinus(), a.numPlus()));
int balance = 0;
int numDaysLeft = k - i;
for (int j = 0; j < n; j++) {
Tickets t = tickets[j];
int col = t.color;
if (t.numMinus() == numDaysLeft) {
int index = t.minusTill;
t.minusTill--;
answer[col][index] = i;
balance--;
} else if (t.numPlus() == numDaysLeft) {
int index = t.plusAfter;
t.plusAfter++;
answer[col][index] = i;
balance++;
} else {
if (t.numMinus() == 0) {
int index = t.plusAfter;
t.plusAfter++;
answer[col][index] = i;
balance++;
} else if (t.numPlus() == 0) {
int index = t.minusTill;
t.minusTill--;
answer[col][index] = i;
balance--;
} else {
if (balance <= 0) {
int index = t.plusAfter;
t.plusAfter++;
answer[col][index] = i;
balance++;
} else {
int index = t.minusTill;
t.minusTill--;
answer[col][index] = i;
balance--;
}
}
}
}
}
grader.allocate_tickets(answer);
long ans = 0;
for (int i = 0; i < n; i++) {
for (int j = 0; j < until[i]; j++) {
ans -= x[i][j];
}
for (int j = after[i]; j < m; j++) {
ans += x[i][j];
}
}
return ans;
}
long doTheJob1(int[][] x, int k, int n, int m) {
Tickets[] tickets = new Tickets[n];
int[][] answer = new int[n][m];
int[] until = new int[n];
// triples [color, index]
PriorityQueue<int[]> pq = new PriorityQueue<int[]>(
(int[] a, int[] b) -> {
int col1 = a[0], i1 = a[1];
int col2 = b[0], i2 = b[1];
return x[col1][i1] != x[col2][i2] ? Integer.compare(x[col1][i1], x[col2][i2]) : col1 - col2;
}
);
for (int i = 0; i < n; i++) {
pq.add(new int[] { i, 0 });
}
for (int i = 0; i < k * (n / 2); i++) {
int[] p = pq.poll();
int col = p[0];
until[col]++;
p[1]++;
if (p[1] < m) {
pq.add(p);
}
}
// int[] after = new int[n];
// PriorityQueue<int[]> pq1 = new PriorityQueue<int[]>(
// (int[] a, int[] b) -> {
// int col1 = a[0], i1 = a[1];
// int col2 = b[0], i2 = b[1];
// return x[col1][i1] != x[col2][i2] ? -Integer.compare(x[col1][i1], x[col2][i2]) : -col1 + col2;
// }
// );
// Arrays.fill(after, m - 1);
// for (int i = 0; i < n; i++) {
// if (until[i] == m) continue;
// pq1.add(new int[] { i, m - 2 });
// }
// for (int i = 0; i < k * (n / 2); i++) {
// int[] p = pq1.poll();
// int col = p[0];
// after[col]--;
// p[1]--;
// if (p[1] >= until[col] - 1) {
// pq1.add(p);
// }
// }
for (int i = 0; i < n; i++) {
//tickets[i] = new Tickets(i, x[i], until[i] - 1, after[i] + 1);
tickets[i] = new Tickets(i, x[i], until[i] - 1, until[i]);
}
for (int i = 0; i < k; i++) {
Arrays.sort(tickets, (Tickets a, Tickets b) -> Math.max(b.numMinus(), b.numPlus()) -
Math.max(a.numMinus(), a.numPlus()));
int balance = 0;
int numDaysLeft = k - i;
for (int j = 0; j < n; j++) {
Tickets t = tickets[j];
int col = t.color;
if (t.numMinus() == numDaysLeft) {
int index = t.minusTill;
t.minusTill--;
answer[col][index] = i;
balance--;
} else if (t.numPlus() == numDaysLeft) {
int index = t.plusAfter;
t.plusAfter++;
answer[col][index] = i;
balance++;
} else {
if (t.numMinus() == 0) {
int index = t.plusAfter;
t.plusAfter++;
answer[col][index] = i;
balance++;
} else if (t.numPlus() == 0) {
int index = t.minusTill;
t.minusTill--;
answer[col][index] = i;
balance--;
} else {
if (balance <= 0) {
int index = t.plusAfter;
t.plusAfter++;
answer[col][index] = i;
balance++;
} else {
int index = t.minusTill;
t.minusTill--;
answer[col][index] = i;
balance--;
}
}
}
}
}
grader.allocate_tickets(answer);
long ans = 0;
for (int i = 0; i < n; i++) {
for (int j = 0; j < until[i]; j++) {
ans -= x[i][j];
}
//for (int j = after[i] + 1; j < m; j++) {
for (int j = until[i]; j < m; j++) {
ans += x[i][j];
}
}
return ans;
}
long find_maximum1(int k, int[][] x) {
int n = x.length;
int m = x[0].length;
if (k <= m) return doTheJob(x, k, n, m);
if (k == m) {
Tickets[] tickets = new Tickets[n];
int[][] answer = new int[n][m];
int[] until = new int[n];
// triples [color, index]
PriorityQueue<int[]> pq = new PriorityQueue<int[]>(
(int[] a, int[] b) -> {
int col1 = a[0], i1 = a[1];
int col2 = b[0], i2 = b[1];
return x[col1][i1] != x[col2][i2] ? Integer.compare(x[col1][i1], x[col2][i2]) : col1 - col2;
}
);
for (int i = 0; i < n; i++) {
pq.add(new int[] { i, 0 });
}
for (int i = 0; i < k * (n / 2); i++) {
int[] p = pq.poll();
int col = p[0];
until[col]++;
p[1]++;
if (p[1] < m) {
pq.add(p);
}
}
for (int i = 0; i < n; i++) {
tickets[i] = new Tickets(i, x[i], until[i] - 1, until[i]);
}
for (int i = 0; i < k; i++) {
Arrays.sort(tickets, (Tickets a, Tickets b) -> Math.max(b.numMinus(), b.numPlus()) -
Math.max(a.numMinus(), a.numPlus()));
int balance = 0;
int numDaysLeft = k - i;
for (int j = 0; j < n; j++) {
Tickets t = tickets[j];
int col = t.color;
if (t.numMinus() == numDaysLeft) {
int index = t.minusTill;
t.minusTill--;
answer[col][index] = i;
balance--;
} else if (t.numPlus() == numDaysLeft) {
int index = t.plusAfter;
t.plusAfter++;
answer[col][index] = i;
balance++;
} else {
if (t.numMinus() == 0) {
int index = t.plusAfter;
t.plusAfter++;
answer[col][index] = i;
balance++;
} else if (t.numPlus() == 0) {
int index = t.minusTill;
t.minusTill--;
answer[col][index] = i;
balance--;
} else {
if (balance <= 0) {
int index = t.plusAfter;
t.plusAfter++;
answer[col][index] = i;
balance++;
} else {
int index = t.minusTill;
t.minusTill--;
answer[col][index] = i;
balance--;
}
}
}
}
}
grader.allocate_tickets(answer);
long ans = 0;
for (int i = 0; i < n; i++) {
for (int j = 0; j < until[i]; j++) {
ans -= x[i][j];
}
for (int j = until[i]; j < m; j++) {
ans += x[i][j];
}
}
return ans;
}
int[][] answer = new int[n][m];
long[][] bal = new long[n][k + 1];
for (int i = 0; i < n; i++) {
bal[i][0] = 0;
for (int j = m - 1; j >= m - k; j--) {
bal[i][0] += x[i][j];
}
for (int j = 1; j <= k; j++) {
bal[i][j] = bal[i][j - 1] - x[i][m - k + j - 1] - x[i][j - 1];
}
}
int nk = n * k;
long[][] dp = new long[n + 1][2 * nk + 1];
for (int i = 0; i <= 2 * nk; i++) {
dp[n][i] = Long.MIN_VALUE;
}
dp[n][nk] = 0;
for (int pos = n - 1; pos >= 0; pos--) {
for (int b = 0; b <= 2 * nk; b++) {
dp[pos][b] = Long.MIN_VALUE;
for (int y = 0; y <= k; y++) {
int newBal = b + y - (k - y);
if (newBal > 2 * nk || newBal < 0) continue;
if (dp[pos + 1][newBal] != Long.MIN_VALUE) {
dp[pos][b] = Math.max(dp[pos][b], dp[pos + 1][newBal] + bal[pos][y]);
}
}
}
}
for (int i = 0; i < n; i++) Arrays.fill(answer[i], -1);
int curBal = nk;
long ans1 = 0;
Tickets[] tickets = new Tickets[n];
for (int i = 0; i < n; i++) {
int take = -1;
for (int y = 0; y <= k; y++) {
int b = curBal + y - (k - y);
if (b < 0 || b > 2 * nk) continue;
if (dp[i + 1][b] != Long.MIN_VALUE && dp[i + 1][b] + bal[i][y] == dp[i][curBal]) {
take = y;
break;
}
}
tickets[i] = new Tickets(i, x[i], take - 1, m - (k - take));
curBal += take - (k - take);
ans1 += bal[i][take];
}
for (int i = 0; i < k; i++) {
Arrays.sort(tickets, (Tickets a, Tickets b) -> Math.max(b.numMinus(), b.numPlus()) -
Math.max(a.numMinus(), a.numPlus()));
int balance = 0;
int numDaysLeft = k - i;
for (int j = 0; j < n; j++) {
Tickets t = tickets[j];
int col = t.color;
if (t.numMinus() == numDaysLeft) {
int index = t.minusTill;
t.minusTill--;
answer[col][index] = i;
balance--;
} else if (t.numPlus() == numDaysLeft) {
int index = t.plusAfter;
t.plusAfter++;
answer[col][index] = i;
balance++;
} else {
if (t.numMinus() == 0) {
int index = t.plusAfter;
t.plusAfter++;
answer[col][index] = i;
balance++;
} else if (t.numPlus() == 0) {
int index = t.minusTill;
t.minusTill--;
answer[col][index] = i;
balance--;
} else {
if (balance <= 0) {
int index = t.plusAfter;
t.plusAfter++;
answer[col][index] = i;
balance++;
} else {
int index = t.minusTill;
t.minusTill--;
answer[col][index] = i;
balance--;
}
}
}
}
}
grader.allocate_tickets(answer);
return ans1;
}
private class Tickets {
int color;
int[] val;
int minusTill;
int plusAfter;
Tickets(int col, int[] val, int minusTill, int plusAfter) {
this.color = col;
this.val = val;
this.minusTill = minusTill;
this.plusAfter = plusAfter;
}
int numMinus() {
return minusTill + 1;
}
int numPlus() {
return val.length - plusAfter;
}
}
}
