#include "rect.h"
#include <utility>
#include <tuple>
#include <vector>
#include <functional>
#include <algorithm>
using namespace std;
using pii = pair<int, int>;
constexpr int MAXN = 2'500;
int N, M;
struct SegTreeMax {
int sz;
vector<pii> tree; // min, index (min(a, b) works)
int ql, qr;
void merge(int i) {
tree[i] = max(tree[2*i], tree[2*i+1]);
}
SegTreeMax(vector<int> V) {
sz = 1;
while (sz < V.size()) sz *= 2;
tree.resize(2*sz);
for (int i = 0; i < V.size(); i++) tree[sz+i] = {V[i], i};
for (int i = V.size(); i < sz; i++) tree[sz+i] = {-1, -1};
for (int i = sz-1; i >= 1; i--) merge(i);
}
pii __query(int ni, int nl, int nr) {
if (qr <= nl || nr <= ql) return {-1, -1};
if (ql <= nl && nr <= qr) return tree[ni];
int nm = (nl+nr)/2;
return max(__query(2*ni, nl, nm), __query(2*ni+1, nm, nr));
}
pii query(int l, int r) {
ql = l;
qr = r;
return __query(1, 0, sz);
}
};
// excluded - excluded, because the last element on both sides is the border
vector<pii> splits(vector<int> V) {
// TODO: pay attention to equal elements. If there are multiples, return the rightmost.
SegTreeMax seg(V);
vector<pii> res;
function<void(int, int, bool, bool)> rec = [&](int l, int r, bool lGood, bool rGood) {
// Exit condition
if (l == r) return;
// Check that the extremes are actual elements
if (l != 0 && r != V.size()) res.push_back({l-1, r});
if (l+1 == r) return;
auto [mx, mxi] = seg.query(l, r); // rightmost max
rec(mxi+1, r, true, rGood);
int lastMxi = mxi;
// iterate from right to left
while (true) {
auto [nxtMx, nxtMxi] = seg.query(l, lastMxi);
if (nxtMx != mx) break;
rec(nxtMxi+1, lastMxi, true, true);
lastMxi = nxtMxi;
}
rec(l, lastMxi, lGood, true);
};
rec(0, V.size(), false, false);
return res;
}
long long count_rectangles(vector<vector<int>> a) {
N = a.size();
M = a[0].size();
vector<vector<int>> a2(M, vector<int>(N));
vector<vector<vector<int>>> grid(N, vector<vector<int>>(M));
for (int i = 0; i < N; i++) {
for (int j = 0; j < M; j++) {
a2[j][i] = a[i][j];
}
}
for (int j = 0; j < M; j++) {
for (auto [s, e]: splits(a2[j])) {
grid[e][j].push_back(s);
}
}
vector<vector<pii>> active(N, vector<pii>(N, {-2, -2})); // (incl, incl)
long long ans = 0;
for (int i = 1; i < N-1; i++) {
vector<pii> x = splits(a[i]);
vector<vector<int>> byEnd(M);
vector<pii> cleanup;
vector<pii> active2(M, {-2, -2});
for (auto[s, e]: x) {
if (active[s][e].second == i-1) {
// extend
active[s][e].second = i;
} else {
// start again
active[s][e] = {i, i};
}
byEnd[e].push_back(s);
}
for (int j = 1; j < M; j++) {
sort(grid[i+1][j-1].rbegin(), grid[i+1][j-1].rend());
sort(byEnd[j].begin(), byEnd[j].end(), [&](int a, int b) {return active[a][j].first > active[b][j].first;});
for (auto e: grid[i+1][j-1]) {
// same idea as above
if (active2[e].second == j-2) active2[e].second = j-1;
else active2[e] = {j-1, j-1};
}
for (int a = 0, b = 0; b < byEnd[j].size(); b++) {
// TODO: might have to use a treap / sparse tree
while (a < grid[i+1][j-1].size() && grid[i+1][j-1][a] >= active[byEnd[j][b]][j].first - 1) {
a++; // ... Insert into treap
}
for (int idx = 0; idx < a; idx++) {
if (active2[grid[i+1][j-1][idx]].first <= byEnd[j][b]+1) {
ans++;
}
}
}
}
}
return ans;
}
