2013. Detect Squares

You are given a stream of points on the X-Y plane. Design an algorithm that:

• Adds new points from the stream into a data structure. Duplicate points are allowed and should be treated as different points.
• Given a query point, counts the number of ways to choose three points from the data structure such that the three points and the query point form an axis-aligned square with positive area.

An axis-aligned square is a square whose edges are all the same length and are either parallel or perpendicular to the x-axis and y-axis.

Implement the DetectSquares class:

• DetectSquares() Initializes the object with an empty data structure.
• void add(int[] point) Adds a new point point = [x, y] to the data structure.
• int count(int[] point) Counts the number of ways to form axis-aligned squares with point point = [x, y] as described above.

class DetectSquares {
    set<int> rows[1001];
    int counts[1001][1001] {0,};

    int getSquares(int x1, int y1, int x2, int y2) {
        if(y1 < 0 || y1 > 1000) return 0;
        return counts[y1][x1] * counts[y1][x2] * counts[y2][x1];
    }
public:
    DetectSquares() {}

    void add(vector<int> point) {
        int y = point[1], x = point[0];
        rows[y].insert(x);
        counts[y][x]++;
    }

    int count(vector<int> point) {
        int y = point[1], x = point[0];
        int res = 0;
        for(auto& otherX : rows[y]) {
            if(otherX == x) continue;
            res += getSquares(otherX, y + abs(x - otherX), x, y) + getSquares(otherX, y - abs(x - otherX), x, y);
        }
        return res;
    }
};

/**
 * Your DetectSquares object will be instantiated and called as such:
 * DetectSquares* obj = new DetectSquares();
 * obj->add(point);
 * int param_2 = obj->count(point);
 */