Detect Squares

Problem Description

You are given a stream of points on a 2D plane. Design an algorithm that:

• Adds new points from the stream to a data structure. Duplicate points are allowed and should be treated as different points.
• Counts the number of ways to form a square with three points already in the data structure such that the query point becomes the fourth corner and the square is axis-aligned (its sides are parallel 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 the query point point = [px, py].

Examples

Constraints

• point.length == 2
• 0 ≤ x, y ≤ 1000
• At most 3000 calls in total will be made to add and count.

Diagonal Match and Corner Frequencies

Let the query point be P = (px, py). To form a square, we can look for any point D = (x, y) that is a diagonal corner to P. For D to be diagonal to P to form an axis-aligned square:

• abs(px - x) == abs(py - y) (equal vertical and horizontal distances, defining a square’s side length d).
• d > 0 (diagonal point cannot be P itself).

If we find a valid diagonal point D = (x, y), the other two corners of the square are uniquely determined:

• C1 = (px, y)
• C2 = (x, py)

The number of squares we can form using this diagonal pair is: frequency(D) * frequency(C1) * frequency(C2)

By storing point counts in a frequency map, we can query C1 and C2 in O(1) time.

Solution: Diagonal Iteration

import java.util.*;

class DetectSquares {
    private int[][] counts = new int[1001][1001]; // coordinates bounded to [0, 1000]
    private List<int[]> pointsList = new ArrayList<>();

    public DetectSquares() {}

    public void add(int[] point) {
        counts[point[0]][point[1]]++;
        pointsList.add(point);
    }

    public int count(int[] point) {
        int px = point[0], py = point[1];
        int sum = 0;

        for (int[] p : pointsList) {
            int x = p[0], y = p[1];

            // check if (x, y) can be a diagonal corner
            if (Math.abs(px - x) == Math.abs(py - y) && px != x) {
                // other two corners are (px, y) and (x, py)
                sum += counts[px][y] * counts[x][py];
            }
        }
        return sum;
    }
}

from collections import defaultdict

class DetectSquares:
    def __init__(self):
        self.counts = defaultdict(int)
        self.points = []

    def add(self, point: list[int]) -> None:
        x, y = point
        self.counts[(x, y)] += 1
        self.points.append((x, y))

    def count(self, point: list[int]) -> int:
        px, py = point
        total = 0

        for x, y in self.points:
            # check if (x, y) is diagonal to (px, py)
            if abs(px - x) == abs(py - y) and px != x:
                # lookup frequency of the other two corners
                total += self.counts[(px, y)] * self.counts[(x, py)]

        return total

#include <vector>
#include <cmath>

class DetectSquares {
    int counts[1001][1001] = {};
    std::vector<std::pair<int, int>> pointsList;

public:
    DetectSquares() {}

    void add(std::vector<int> point) {
        counts[point[0]][point[1]]++;
        pointsList.push_back({point[0], point[1]});
    }

    int count(std::vector<int> point) {
        int px = point[0], py = point[1];
        int sum = 0;

        for (const auto& [x, y] : pointsList) {
            if (std::abs(px - x) == std::abs(py - y) && px != x) {
                sum += counts[px][y] * counts[x][py];
            }
        }
        return sum;
    }
};

Complexity Analysis:

• add: O(1).
• count: O(N) where N is the number of points added. We iterate over all points.
• Space Complexity: O(N) to store points, plus O(1001 * 1001) grid space for the frequency lookup matrix.

← All Problems