View on NeetCode
View on LeetCode

Problem

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.

Example 1:

Input
["DetectSquares", "add", "add", "add", "count", "count", "add", "count"]
[[], [[3, 10]], [[11, 2]], [[3, 2]], [[11, 10]], [[14, 8]], [[11, 2]], [[11, 10]]]
Output
[null, null, null, null, 1, 0, null, 2]

Explanation
DetectSquares detectSquares = new DetectSquares();
detectSquares.add([3, 10]);
detectSquares.add([11, 2]);
detectSquares.add([3, 2]);
detectSquares.count([11, 10]); // return 1. You can choose:
                                //   - The first, second, and third points
detectSquares.count([14, 8]);  // return 0. The query point cannot form a square with any points in the data structure.
detectSquares.add([11, 2]);     // Adding duplicate points is allowed.
detectSquares.count([11, 10]); // return 2. You can choose:
                                //   - The first, second, and third points
                                //   - The first, third, and fourth points

Constraints:

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

Solution

Approach: Hash Map with Point Counting

The key insight is to store point counts in a hash map, then for each query point, iterate through all possible diagonal corners to find squares.

Implementation

from collections import defaultdict

class DetectSquares:
    def __init__(self):
        # Store count of each point
        self.point_count = defaultdict(int)
        # Store all unique points for iteration
        self.points = []

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

    def count(self, point: List[int]) -> int:
        qx, qy = point
        total = 0

        # For each point, check if it can be a diagonal corner
        for px, py in self.points:
            # Skip if same point or not diagonal
            if abs(px - qx) != abs(py - qy) or px == qx or py == qy:
                continue

            # Found a valid diagonal corner
            # Check if the other two corners exist
            count1 = self.point_count[(px, qy)]
            count2 = self.point_count[(qx, py)]

            # Multiply counts (for duplicate points)
            total += count1 * count2

        return total

Approach 2: Optimized with Points by X-coordinate

from collections import defaultdict

class DetectSquares:
    def __init__(self):
        # Map from x-coordinate to list of y-coordinates
        self.x_to_points = defaultdict(list)
        # Count of each point
        self.point_count = defaultdict(int)

    def add(self, point: List[int]) -> None:
        x, y = point
        self.x_to_points[x].append(y)
        self.point_count[(x, y)] += 1

    def count(self, point: List[int]) -> int:
        qx, qy = point
        total = 0

        # Iterate through points with same x-coordinate as query
        if qx not in self.x_to_points:
            return 0

        for py in self.x_to_points[qx]:
            if py == qy:
                continue

            side_length = abs(qy - py)

            # Check both possible x-coordinates for the square
            for dx in [side_length, -side_length]:
                diagonal_x = qx + dx

                # Count squares with this diagonal
                count1 = self.point_count[(diagonal_x, qy)]
                count2 = self.point_count[(diagonal_x, py)]
                total += count1 * count2

        return total

Approach 3: Using Tuples for Point Storage

from collections import Counter

class DetectSquares:
    def __init__(self):
        self.points = Counter()

    def add(self, point: List[int]) -> None:
        self.points[tuple(point)] += 1

    def count(self, point: List[int]) -> int:
        qx, qy = point
        total = 0

        for (px, py), count in self.points.items():
            # Check if this point can be diagonal corner
            if abs(px - qx) != abs(py - qy) or px == qx or py == qy:
                continue

            # Count the two other corners
            total += count * self.points[(px, qy)] * self.points[(qx, py)]

        return total

Complexity Analysis

  • add() Time Complexity: O(1), just adding to hash map and list.
  • count() Time Complexity: O(n), where n is the number of points. We check each point as a potential diagonal corner.
  • Space Complexity: O(n) for storing all points.

Key Insights

  1. Diagonal Corner Strategy: For any square, fix the query point and one diagonal corner, then check if the other two corners exist.

  2. Square Properties: An axis-aligned square has:
    • All sides equal length
    • Diagonal corners satisfy x1 - x2 = y1 - y2
  3. Duplicate Handling: Using counts allows us to handle duplicate points by multiplication.

  4. Point Counting: Store each point with its count to handle duplicates efficiently.

  5. Three Points Needed: Given query point, we need to find 3 other points. Fix one diagonal corner and check if the other 2 exist.

  6. Optimization: Can optimize by grouping points by x-coordinate or y-coordinate to reduce search space.

  7. Validation: For diagonal corners (qx, qy) and (px, py):
    • Must satisfy: px - qx = py - qy (equal sides)
    • Must satisfy: px ≠ qx and py ≠ qy (not same point, not on same axis)