Detect Squares

Problem

https://leetcode.com/problems/detect-squares/

You are given a stream of points on the X-Y plane. Design a data structure 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.

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]

Constraints:

  • point.length == 2

  • 0 <= x, y <= 1000

  • At most 3000 calls in total will be made to add and count.

Pattern

Hash Table, Math

Approaches

Code

from collections import defaultdict


class DetectSquares:
    """Detect axis-aligned squares from a stream of points."""

    def __init__(self) -> None:
        self.points: dict[tuple[int, int], int] = defaultdict(int)
        self.points_list: list[tuple[int, int]] = []

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

    def count(self, point: list[int]) -> int:
        res = 0
        qx, qy = point
        for x, y in self.points_list:
            if abs(qx - x) == abs(qy - y) and qx != x and qy != y:
                res += self.points[(qx, y)] * self.points[(x, qy)]
        return res

Test

>>> from detect_squares__hash_map import DetectSquares
>>> ds = DetectSquares()
>>> ds.add([3, 10])
>>> ds.add([11, 2])
>>> ds.add([3, 2])
>>> ds.count([11, 10])
1
>>> ds.count([14, 8])
0
>>> ds.add([11, 2])
>>> ds.count([11, 10])
2
class detect_squares__hash_map.DetectSquares

Bases: object

Detect axis-aligned squares from a stream of points.

add(point: list[int]) None
count(point: list[int]) int