#include "fish.h"
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
ll max_weights(int N, int M, vector<int> X, vector<int> Y, vector<int> W) {
vector<vector<int>> fish(N);
for (int i = 0; i < M; i++) {
fish[X[i]].push_back(i);
}
vector<vector<ll>> pfs(N);
for (int i = 0; i < N; i++) {
sort(fish[i].begin(), fish[i].end(), [&](int x, int y) {return Y[x] < Y[y];});
pfs[i].resize(fish[i].size() + 1);
for (int j = 0; j < fish[i].size(); j++) {
pfs[i][j + 1] = pfs[i][j] + W[fish[i][j]];
}
}
auto get_sum = [&](int c, int lo, int hi) {
if (c == -1 || c >= N) {
return 0LL;
}
int x = lower_bound(fish[c].begin(), fish[c].end(), lo, [&](int i, int t) {return Y[i] < t;}) - fish[c].begin();
int y = lower_bound(fish[c].begin(), fish[c].end(), hi, [&](int i, int t) {return Y[i] < t;}) - fish[c].begin();
return pfs[c][y] - pfs[c][x];
};
vector<array<ll, 2>> dp{{0, 0}};
vector<int> h{-1};
for (int i = 0; i < N; i++) {
vector<int> cand{-1};
if (i) {
for (int j = 0; j < fish[i - 1].size(); j++) {
cand.push_back(Y[fish[i - 1][j]]);
}
}
if (i < N - 1) {
for (int j = 0; j < fish[i + 1].size(); j++) {
cand.push_back(Y[fish[i + 1][j]]);
}
}
sort(cand.begin(), cand.end());
vector pd(cand.size(), array<ll, 2>{0, 0});
for (int j = 0; j < h.size(); j++) {
pd[0][0] = max(pd[0][0], dp[j][1]);
pd[0][1] = max(pd[0][1], dp[j][1] + get_sum(i, 0, h[j] + 1));
}
for (int j = 1; j < cand.size(); j++) {
for (int k = 0; k < h.size(); k++) {
if (h[k] <= cand[j]) {
pd[j][0] = max(pd[j][0], dp[k][0] + get_sum(i - 1, h[k] + 1, cand[j] + 1));
} else {
pd[j][1] = max(pd[j][1], dp[k][1] + get_sum(i, cand[j] + 1, h[k] + 1));
}
pd[j][0] = max(pd[j][0], dp[0][1]);
pd[j][1] = max(pd[j][1], pd[j][0]);
}
}
swap(dp, pd);
swap(h, cand);
}
ll ans = 0;
for (auto &v : dp) {
ans = max(ans, v[1]);
}
return ans;
}
