Problem statement

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

Solution

The idea is to keep two pieces of infromation:

  1. self.d is counter for points, that is how many times we have each of them.
  2. self.x_coord is defaultdict of counter for each coordinate x, that is given x we can quickly find all point with this coordinate.

What we need to do now:

  1. for add(self, point) we just increment two counters.
  2. for count(self, point) we need to find all points with the same coordinate x: points in the form (x, y2) and then reconstruct square: there will be two ways to do it: one above line and one below. Here we need to take into account frequences of our points, so we use self.d for this. Be careful to check that y2 != y, it costs me 5 minutes penalty on the contest.

Complexity

Time complexity of add is O(1). Time complexity of count is potentially O(n) but in practice it is less because usually we do not have a lot of points on the same line. Space complexity of all data structure is O(n).

Code

class DetectSquares:
    def __init__(self):
        self.d = Counter()
        self.x_coord = defaultdict(Counter)   # x -> all y coordinates with freqs

    def add(self, point):
        x, y = point
        self.d[x, y] += 1
        self.x_coord[x][y] += 1

    def count(self, point):
        x, y = point
        ans = 0
        for y2 in self.x_coord[x]:
            if y == y2: continue
            ans += self.d[x, y2] * self.d[x + y2 - y, y] * self.d[x + y2 - y, y2]
            ans += self.d[x, y2] * self.d[x + y - y2, y] * self.d[x + y - y2, y2]
        return ans