Problem Description

You need to design a data structure that can store points on a 2D plane and count how many axis-aligned squares can be formed using stored points and a query point.

An axis-aligned square is a square where:

• All four sides have the same length
• All sides are either parallel or perpendicular to the x-axis and y-axis (no rotation)
• The square has positive area (not a single point)

The DetectSquares class should support three operations:

1. Constructor DetectSquares(): Initializes an empty data structure to store points.

2. Method add(point): Takes a point [x, y] and adds it to the data structure. Duplicate points are allowed - if you add the same point multiple times, each occurrence counts as a separate point.

3. Method count(point): Takes a query point [x, y] and returns the number of ways to choose three points from the stored data structure such that these three points together with the query point form an axis-aligned square.

For the count method, you're essentially finding all possible combinations of three stored points that, when combined with the query point as the fourth corner, create a valid axis-aligned square. Since duplicate points are treated as different points, if a point appears multiple times in the data structure, each occurrence can be used separately in forming different squares.

For example, if you have points at (0,0), (0,1), (1,0), and (1,1) in your data structure, and you query with point (0,0), you would count how many ways you can pick three points from the stored points to form a square with (0,0) as one corner.

Intuition

When we need to form an axis-aligned square with a query point, we need to think about what constraints define such a square. Since the square must be axis-aligned, its sides must be parallel to the x and y axes. This means if we have one corner at (x1, y1), the other three corners must follow a specific pattern.

The key insight is that for any axis-aligned square, if we know two diagonal corners, we can determine the other two corners. Given a query point (x1, y1) as one corner, we can think of it forming a diagonal with another point. But which diagonal point should we consider?

Instead of thinking about diagonals directly, let's think about the structure of an axis-aligned square. If (x1, y1) is one corner, and we want to form a square with side length d, then:

• One adjacent corner could be at (x2, y1) where |x2 - x1| = d (horizontally aligned)
• From these two points, the remaining two corners must be at (x1, y1 ± d) and (x2, y1 ± d)

This gives us a systematic way to search: for each point that shares the same y-coordinate as our query point but has a different x-coordinate, we can calculate the distance d = |x2 - x1|. This distance becomes our potential square side length. Then we just need to check if the other two required corners exist in our data structure.

Since we might have the square extending either upward or downward from the horizontal edge we found, we need to check both possibilities:

• Square above: corners at (x1, y1 + d) and (x2, y1 + d)
• Square below: corners at (x1, y1 - d) and (x2, y1 - d)

For storage, a nested hash map structure cnt[x][y] works perfectly because:

• We can quickly check if a specific point exists
• We can track the count of duplicate points (important since duplicates are treated as separate points)
• We can easily iterate through all x-coordinates to find potential horizontal partners

The counting logic uses multiplication because if point A appears m times, point B appears n times, and point C appears p times, then the total number of ways to pick one of each is m × n × p.

Solution Approach

The solution uses a nested hash table structure to efficiently store and query points. Here's how each component works:

Data Structure:

• self.cnt = defaultdict(Counter) creates a nested dictionary where:
  • The outer dictionary keys are x-coordinates
  • The inner Counter objects map y-coordinates to their occurrence counts
  • This structure allows cnt[x][y] to give us the count of point (x, y)

Add Method:

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

• Simply increments the count for point (x, y)
• Handles duplicates naturally by maintaining a count

Count Method:

def count(self, point: List[int]) -> int:
    x1, y1 = point
    if x1 not in self.cnt:
        return 0
    ans = 0
    for x2 in self.cnt.keys():
        if x2 != x1:
            d = x2 - x1
            ans += self.cnt[x2][y1] * self.cnt[x1][y1 + d] * self.cnt[x2][y1 + d]
            ans += self.cnt[x2][y1] * self.cnt[x1][y1 - d] * self.cnt[x2][y1 - d]
    return ans

The counting algorithm works as follows:

1. Early termination: If there are no points with x-coordinate x1, we can't form any squares, so return 0.

2. Enumerate horizontal partners: For each x-coordinate x2 in our data structure where x2 ≠ x1:

   • Calculate the distance d = x2 - x1 (this will be the side length of potential squares)
   • Note: d can be positive or negative depending on whether x2 is to the right or left of x1

3. Check for two possible squares:

   • Square extending upward: We need points at:
     • (x2, y1) - the horizontal partner
     • (x1, y1 + d) - vertical from query point
     • (x2, y1 + d) - the diagonal corner
   • Square extending downward: We need points at:
     • (x2, y1) - the horizontal partner
     • (x1, y1 - d) - vertical from query point
     • (x2, y1 - d) - the diagonal corner

4. Count combinations: For each valid square configuration, multiply the counts of all three required points. This gives us the number of ways to choose one instance of each point when duplicates exist.

The beauty of this approach is that by fixing the horizontal edge first (from (x1, y1) to (x2, y1)), we deterministically know where the other two corners must be for an axis-aligned square. We don't need to search through all possible combinations of three points - we only need to verify if the required points exist at the calculated positions.

Example Walkthrough

Let's walk through a concrete example to illustrate how the solution works.

Setup: Suppose we've added the following points to our data structure:

• add([0, 0])
• add([0, 2])
• add([2, 0])
• add([2, 2]) - added twice (duplicate)
• add([2, 2])
• add([1, 1])

Our data structure cnt now looks like:

cnt[0][0] = 1
cnt[0][2] = 1
cnt[2][0] = 1
cnt[2][2] = 2  (appears twice)
cnt[1][1] = 1

Query: count([0, 0])

Let's count how many squares can be formed with (0, 0) as one corner.

1. Extract query point: x1 = 0, y1 = 0

2. Check if x1 exists: x1 = 0 is in cnt, so we continue.

3. Iterate through all x-coordinates:

   When x2 = 0: Skip since x2 == x1

   When x2 = 2:

   • Calculate distance: d = 2 - 0 = 2
   • Check upward square (need points at (2,0), (0,2), (2,2)):
     • cnt[2][0] = 1 ✓
     • cnt[0][0 + 2] = cnt[0][2] = 1 ✓
     • cnt[2][0 + 2] = cnt[2][2] = 2 ✓
     • Count: 1 × 1 × 2 = 2 squares
   • Check downward square (need points at (2,0), (0,-2), (2,-2)):
     • cnt[2][0] = 1 ✓
     • cnt[0][0 - 2] = cnt[0][-2] = 0 ✗ (doesn't exist)
     • Count: 1 × 0 × 0 = 0 squares

   When x2 = 1:

   • Calculate distance: d = 1 - 0 = 1
   • Check upward square (need points at (1,0), (0,1), (1,1)):
     • cnt[1][0] = 0 ✗ (doesn't exist)
     • Count: 0 × ... = 0 squares
   • Check downward square (need points at (1,0), (0,-1), (1,-1)):
     • cnt[1][0] = 0 ✗ (doesn't exist)
     • Count: 0 × ... = 0 squares

4. Total count: 2 + 0 + 0 = 2

The function returns 2, meaning there are 2 ways to form squares with (0,0) as one corner. Both ways use the points (0,0), (2,0), (0,2), and one instance of (2,2). Since (2,2) appears twice in our data structure, we can choose either occurrence, giving us 2 different combinations.

Visual representation of the found square:

(0,2) ---- (2,2)
  |          |
  |          |
(0,0) ---- (2,0)

This example demonstrates how the algorithm efficiently finds squares by:

1. Fixing a horizontal edge from the query point
2. Calculating the required positions of the other two corners
3. Multiplying the counts to handle duplicates correctly

Solution Implementation

Time and Space Complexity

Time Complexity:

• add(point): The operation involves accessing a nested dictionary and incrementing a counter value. Both dictionary access self.cnt[x] and counter increment self.cnt[x][y] += 1 are O(1) operations. Therefore, the time complexity is O(1).

• count(point): The method iterates through all unique x-coordinates in self.cnt.keys(). For each x-coordinate x2 different from x1, it performs constant-time lookups in the nested dictionary structure (accessing counter values). The number of unique x-coordinates can be at most n where n is the total number of points added. Therefore, the time complexity is O(n).

Space Complexity:

• The data structure uses a nested dictionary where the outer dictionary stores x-coordinates as keys, and each value is a Counter object storing y-coordinate frequencies. In the worst case, all n points added could be unique, requiring storage for all of them. Therefore, the space complexity is O(n), where n is the number of points in the data stream.

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Attempting to Find All Four Corners Instead of Three

Pitfall: A common mistake is trying to find all four corners of a square when one corner (the query point) is already given. This leads to overcounting since the query point shouldn't be selected from the stored points.

Incorrect approach:

def count(self, point: List[int]) -> int:
    x1, y1 = point
    ans = 0
    # Wrong: Looking for all 4 corners including the query point
    for x2 in self.point_count:
        if x2 != x1:
            d = x2 - x1
            # This would count the query point if it exists in storage
            ans += (self.point_count[x1][y1] *  # WRONG! Don't count query point
                   self.point_count[x2][y1] * 
                   self.point_count[x1][y1 + d] * 
                   self.point_count[x2][y1 + d])
    return ans

Solution: The query point serves as one fixed corner. You only need to find three other points from the stored data structure to form a square.

2. Forgetting to Check Both Upward and Downward Squares

Pitfall: When fixing a horizontal edge from (x1, y1) to (x2, y1), there are two possible squares: one extending upward and one extending downward. Many implementations only check one direction.

Incorrect approach:

def count(self, point: List[int]) -> int:
    x1, y1 = point
    ans = 0
    for x2 in self.point_count:
        if x2 != x1:
            d = abs(x2 - x1)  # Using abs loses direction information
            # Only checking upward squares
            ans += (self.point_count[x2][y1] * 
                   self.point_count[x1][y1 + d] * 
                   self.point_count[x2][y1 + d])
    return ans

Solution: Always check both y1 + d and y1 - d directions. The side length d should maintain its sign (positive or negative) to ensure the square is properly formed.

3. Using Absolute Value for Distance Calculation

Pitfall: Using abs(x2 - x1) for the side length breaks the geometric relationship needed to form a proper square.

Why it's wrong: If x2 is to the left of x1 (negative d), the square extending "upward" actually needs points at y1 + d (which would be y1 + negative value = y1 - |d|). Using absolute value would incorrectly look for points at y1 + |d|.

Correct approach: Keep the signed distance d = x2 - x1. This ensures:

• When x2 > x1 (d > 0): Forms squares to the right
• When x2 < x1 (d < 0): Forms squares to the left The signed distance maintains the correct geometric relationships for both cases.

4. Not Handling Duplicate Points Correctly

Pitfall: Treating duplicate points as a single point or not properly counting all combinations when duplicates exist.

Example scenario: If point (0, 1) is added 3 times and forms part of a valid square, you should count 3 different ways to pick this point.

Solution: Use a Counter or dictionary to track the count of each point, then multiply the counts when calculating combinations. The formula cnt[x2][y1] * cnt[x1][y1 + d] * cnt[x2][y1 + d] naturally handles this by multiplication.

5. Inefficient Point Storage Structure

Pitfall: Using a list of points and checking all combinations with nested loops leads to O(n³) complexity for the count operation.

Inefficient approach:

def count(self, point: List[int]) -> int:
    points = self.stored_points  # List of all points
    count = 0
    # O(n³) - checking all triplets
    for p1 in points:
        for p2 in points:
            for p3 in points:
                if forms_square(point, p1, p2, p3):
                    count += 1
    return count

Solution: Use the nested hash table structure defaultdict(Counter) for O(1) point lookup and O(n) count operation where n is the number of unique x-coordinates.