IOI '19 - Rectangles

#include "rect.h"
#include <iostream>
#include <stack>
#include <tuple>

using namespace std;

enum direction { Left, Right, Up, Down };
int dr[]{0, 0, -1, +1}, dc[]{-1, +1, 0, 0};

const int N = 2500, LG = 12;
int maxP2[1+N]{};

long long count_rectangles(vector<vector<int>> a) {
    int n = a.size(), m = a[0].size();

    vector<vector<int>> reach[4]; // distance to closest in each direction that is greater than or equal to
    for (int d = 0; d < 4; ++d) reach[d].assign(n, vector<int>(m));
    for (int i = 0; i < n; ++i) { // row
        { // left
            stack<int> candidates;
            for (int j = 0; j < m; ++j) {
                while (!candidates.empty() && a[i][candidates.top()] < a[i][j]) candidates.pop();
                reach[0][i][j] = candidates.empty() ? -1 : candidates.top();
                reach[0][i][j] = j - reach[0][i][j];
                candidates.push(j);
            }
        }
        { // right
            stack<int> candidates;
            for (int j = m-1; j >= 0; --j) {
                while (!candidates.empty() && a[i][candidates.top()] < a[i][j]) candidates.pop();
                reach[1][i][j] = candidates.empty() ? m : candidates.top();
                reach[1][i][j] = reach[1][i][j] - j;
                candidates.push(j);
            }
        }
    }
    for (int j = 0; j < m; ++j) { // column
        { // up
            stack<int> candidates;
            for (int i = 0; i < n; ++i) {
                while (!candidates.empty() && a[candidates.top()][j] < a[i][j]) candidates.pop();
                reach[2][i][j] = candidates.empty() ? -1 : candidates.top();
                reach[2][i][j] = i - reach[2][i][j];
                candidates.push(i);
            }
        }
        { // down
            stack<int> candidates;
            for (int i = n-1; i >= 0; --i) {
                while (!candidates.empty() && a[candidates.top()][j] < a[i][j]) candidates.pop();
                reach[3][i][j] = candidates.empty() ? n : candidates.top();
                reach[3][i][j] = reach[3][i][j] - i;
                candidates.push(i);
            }
        }
    }

    /* The letters below represent values just outside the border of the rectangle, adjacent to corners of the rectangle
         a     e
        .-------.
       c|       |d
        |       |
        |       |
        .-------.
         b     f
    */

    vector<tuple<int, int, int, int>> rectangles;
    auto consider = [&](int i, int j, int h, int w) { // rectangle with given top left position, width, height
        if (i <= 0 || j <= 0 || i+h-1 >= n-1 || j+w-1 >= m-1) return;
        if ( w < 1 || h < 1) return;
        // check that all sides can reach one another
        // if (RMQ(Right, j-1, i, i+h-1) <= w) return;
        // if (RMQ(Left, j+w, i, i+h-1) <= w) return;
        // if (RMQ(Down, i-1, j, j+w-1) <= h) return;
        // if (RMQ(Up, i+h, j, j+w-1) <= h) return;
        // ++rectangles;
        rectangles.emplace_back(i, j, h, w);
    };
    for (int i = 1; i < n-1; ++i) {
        for (int j = 1; j < m-1; ++j) {
            { // Case 1: a <= b && c <= d: consider i, j as top left
                int height = reach[Down][i-1][j] - 1, width = reach[Right][i][j-1] - 1;
                consider(i, j, height, width);
            }

            { // Case 2: a <= b && c > d: consider i, j as top right
                int width = reach[Left][i][j+1] - 1;
                int jLeft = j-width+1;
                if (jLeft > 0 && a[i][jLeft-1] > a[i][j+1]) { // c > d, not equal
                    int height = reach[Down][i-1][jLeft] - 1;
                    consider(i, jLeft, height, width);
                }
            }

            { // Case 3: a > b && c <= d: consider i, j as bottom left
                int height = reach[Up][i+1][j] - 1;
                int iTop = i-height+1;
                if (iTop > 0 && a[iTop-1][j] > a[i+1][j]) { // a > b, not equal
                    int width = reach[Right][iTop][j-1] - 1;
                    consider(iTop, j, height, width);
                }
            }

            { // Case 4: a > b && c > d && e <= f: consider i, j, as top right
                int height = reach[Down][i-1][j] - 1, width = reach[Left][i][j+1] - 1;
                int iBottom = i+height-1, jLeft = j-width+1;
                if (iBottom+1 < n && jLeft > 0) {
                    // check that a > b and c > d
                    if (a[i-1][jLeft] > a[iBottom+1][jLeft] && a[i][jLeft-1] > a[i][j+1]) {
                        consider(i, jLeft, height, width);
                    }
                }
            }

            { // Case 5: a > b && c > d && e > f: consider i, j as bottom right
                int height = reach[Up][i+1][j] - 1;
                int iTop = i-height+1;
                if (iTop > 0 && a[iTop-1][j] > a[i+1][j]) { // e > f, not equal
                    int width = reach[Left][iTop][j+1] - 1;
                    int jLeft = j-width+1;
                    // check that a > b and c > d
                    if (a[iTop-1][jLeft] > a[i+1][jLeft] && a[iTop][jLeft-1] > a[iTop][j+1]) {
                        consider(iTop, jLeft, height, width);
                    }
                }
            }
        }
    }

    // test reach in all four directions separately to reduce memory usage by a factor of 4

    // maxP2
    for (int i = 2; i <= max(n, m); ++i) maxP2[i] = maxP2[i-1] + !(i & (i-1));

    vector<vector<int>> reachRMQ[1+LG];
    for (int k = 0; k <= LG; ++k) reachRMQ[k].assign(n, vector<int>(m));
    auto RMQ = [&](int dir, int rc, int l, int r) {
        int p2 = maxP2[r-l+1];
        if (dir == 0 || dir == 1) {
            int j = rc;
            return min(reachRMQ[p2][l][j], reachRMQ[p2][r-(1<<p2)+1][j]);
        } else {
            int i = rc;
            return min(reachRMQ[p2][i][l], reachRMQ[p2][i][r-(1<<p2)+1]);
        }
    };

    vector<tuple<int, int, int, int>> rectangles2;

    // Left
    for (int j = 0; j < m; ++j) { // column
        for (int i = 0; i < n; ++i) reachRMQ[0][i][j] = reach[0][i][j];
        for (int k = 1; k <= LG; ++k) {
            for (int i = 0; i < n; ++i) {
                int i2 = i + (1 << (k-1));
                if (i2 >= n) {
                    reachRMQ[k][i][j] = reachRMQ[k-1][i][j];
                } else {
                    reachRMQ[k][i][j] = min(reachRMQ[k-1][i][j], reachRMQ[k-1][i + (1 << (k-1))][j]);
                }
            }
        }
    }
    rectangles2.clear();
    for (auto [i, j, h, w] : rectangles) {
        if (RMQ(Left, j+w, i, i+h-1) > w) rectangles2.emplace_back(i, j, h, w);
    }
    rectangles = rectangles2;

    // Right
    for (int k = 0; k <= LG; ++k) reachRMQ[k].assign(n, vector<int>(m));
    for (int j = 0; j < m; ++j) { // column
        for (int i = 0; i < n; ++i) reachRMQ[0][i][j] = reach[1][i][j];
        for (int k = 1; k <= LG; ++k) {
            for (int i = 0; i < n; ++i) {
                int i2 = i + (1 << (k-1));
                if (i2 >= n) {
                    reachRMQ[k][i][j] = reachRMQ[k-1][i][j];
                } else {
                    reachRMQ[k][i][j] = min(reachRMQ[k-1][i][j], reachRMQ[k-1][i + (1 << (k-1))][j]);
                }
            }
        }
    }
    rectangles2.clear();
    for (auto [i, j, h, w] : rectangles) {
        if (RMQ(Right, j-1, i, i+h-1) > w) rectangles2.emplace_back(i, j, h, w);
    }
    rectangles = rectangles2;

    // Up
    for (int i = 0; i < n; ++i) { // row
        for (int j = 0; j < m; ++j) reachRMQ[0][i][j] = reach[2][i][j];
        for (int k = 1; k <= LG; ++k) {
            for (int j = 0; j < m; ++j) {
                int j2 = j + (1 << (k-1));
                if (j2 >= m) {
                    reachRMQ[k][i][j] = reachRMQ[k-1][i][j];
                } else {
                    reachRMQ[k][i][j] = min(reachRMQ[k-1][i][j], reachRMQ[k-1][i][j + (1 << (k-1))]);
                }
            }
        }
    }
    rectangles2.clear();
    for (auto [i, j, h, w] : rectangles) {
        if (RMQ(Up, i+h, j, j+w-1) > h) rectangles2.emplace_back(i, j, h, w);
    }
    rectangles = rectangles2;

    // Down
    for (int i = 0; i < n; ++i) { // row
        for (int j = 0; j < m; ++j) reachRMQ[0][i][j] = reach[3][i][j];
        for (int k = 1; k <= LG; ++k) {
            for (int j = 0; j < m; ++j) {
                int j2 = j + (1 << (k-1));
                if (j2 >= m) {
                    reachRMQ[k][i][j] = reachRMQ[k-1][i][j];
                } else {
                    reachRMQ[k][i][j] = min(reachRMQ[k-1][i][j], reachRMQ[k-1][i][j + (1 << (k-1))]);
                }
            }
        }
    }
    rectangles2.clear();
    for (auto [i, j, h, w] : rectangles) {
        if (RMQ(Down, i-1, j, j+w-1) > h) rectangles2.emplace_back(i, j, h, w);
    }
    rectangles = rectangles2;

    return rectangles.size();
}