#include "fish.h"
#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
using ll = long long;
const ll INF = 1e18;
ll max_weights(int N, int M, vector<int> X, vector<int> Y, vector<int> W) {
vector<vector<pair<int, ll>>> fish(N+1); // fish[col] = { {row, weight}, ... }
vector<vector<int>> anc(N+1); // anc[col] = { row, ... }
for (int i = 0; i < M; i++) {
int row = Y[i]+1, col = X[i]+1, wgt = W[i];
fish[col].push_back({row, wgt});
if (col > 1) anc[col-1].push_back(row);
anc[col].push_back(row);
if (col < N) anc[col+1].push_back(row);
}
// sort ancs
for (int i = 0; i <= N; i++) {
anc[i].push_back(0);
sort(anc[i].begin(), anc[i].end());
anc[i].erase(unique(anc[i].begin(), anc[i].end()), anc[i].end());
// cout << "anc " << i << " = ";
// for (int j : anc[i]) cout << j << ' ';
// cout << endl;
}
// sort + prefsum fish
for (int i = 0; i <= N; i++) {
fish[i].push_back({0, 0});
sort(fish[i].begin(), fish[i].end());
ll sum = 0;
for (auto &[_, wgt] : fish[i]) {
sum += wgt;
wgt = sum;
}
// cout << "fish " << i << " = ";
// for (auto [j, w] : fish[i]) cout << j << ':' << w << ' ';
// cout << endl;
}
auto cost = [&] (int i, int j) {
int row = lower_bound(
fish[i].begin(),
fish[i].end(),
pair<int, ll>{j+1, 0}) - fish[i].begin() - 1;
return fish[i][row].second;
};
vector<vector<ll>> inc(N+1), dec(N+1);
inc[0] = {0};
dec[0] = {0};
for (int i = 1; i <= N; i++) {
// cout << i << " ========" << endl;
inc[i].resize(anc[i].size());
dec[i].resize(anc[i].size());
// inc[i-1] -> inc[i]
ll best = -INF;
for (int j = 0, k = 0; j < (int)anc[i].size(); j++) {
while (k < (int)anc[i-1].size() && anc[i-1][k] <= anc[i][j]) {
best = max(best, inc[i-1][k] - cost(i-1, anc[i-1][k]));
k++;
}
inc[i][j] = max(inc[i][j], best + cost(i-1, anc[i][j]));
}
// dec[i-1] -> dec[i]
best = -INF;
for (int j = 0, k = anc[i-1].size()-1; j < (int)anc[i].size(); j++) {
while (k >= 0 && anc[i-1][k] >= anc[i][j]) {
best = max(best, dec[i-1][k] + cost(i, anc[i-1][k]));
k--;
}
dec[i][j] = max(dec[i][j], best - cost(i, anc[i][j]));
}
// dec[i-2] -> inc[i]
if (i-2 >= 0) {
best = -INF;
for (int j = 0, k = 0; j < (int)anc[i].size(); j++) {
while (k < (int)anc[i-2].size() && anc[i-2][k] <= anc[i][j]) {
best = max(best, dec[i-2][k]);
k++;
}
inc[i][j] = max(inc[i][j], best + cost(i-1, anc[i][j]));
}
best = -INF;
for (int j = 0, k = anc[i-2].size()-1; j < (int)anc[i].size(); j++) {
while (k >= 0 && anc[i-2][k] >= anc[i][j]) {
best = max(best, dec[i-2][k] + cost(i-1, anc[i-2][k]));
k--;
}
dec[i][j] = max(dec[i][j], best);
}
}
// inc[i] -> dec[i]
for (int j = 0; j < (int)anc[j].size(); j++) {
dec[i][j] = max(dec[i][j], inc[i][j]);
// cout << j << " -> " << inc[i][j] << ' ' << dec[i][j] << endl;
}
}
ll res = 0;
for (ll x : dec[N]) res = max(res, x);
return res;
}
