Leetcode 2013. Detect Squares Take of care of dups

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.

---

Wrong code with rectangles instead of squares:

from collections import defaultdict
from typing import List

class DetectSquares:

    def __init__(self):
        self.x_to_points_dic = defaultdict(list)
        self.y_to_points_dic = defaultdict(list)
        self.points_cnt_dic = defaultdict(int)

    def add(self, point: List[int]) -> None:
        self.x_to_points_dic[point[0]].append(point)
        self.y_to_points_dic[point[1]].append(point)
        self.points_cnt_dic[tuple(point)] += 1
        
    def count(self, point: List[int]) -> int:
        res = 0
        x, y = point[0], point[1]
        print('point={} self.x_to_points_dic[x]={} self.y_to_points_dic[y]={}'.format(point, self.x_to_points_dic[x], self.y_to_points_dic[y]))
        for px in self.x_to_points_dic[x]:
            for py in self.y_to_points_dic[y]:
                if (py[0],px[1]) in self.points_cnt_dic and (py[0] != x) and (px[1] != y):
                    print('new_point={}'.format((py[0],px[1])))
                    res += self.points_cnt_dic[(py[0],px[1])]
        return res


# Your DetectSquares object will be instantiated and called as such:
obj = DetectSquares()
obj.add([5,10]),
obj.add([10,5]),
obj.add([10,10]),
# obj.count([5,5]),
obj.add([3,0]),
obj.add([8,0]),
obj.add([8,5]),
# obj.count([3,5]),
obj.add([9,0]),
obj.add([9,8]),
obj.add([1,8]),
# obj.count([1,0]),
obj.add([0,0]),
obj.add([8,0]),
obj.add([8,8]),
print(obj.count([0,8]))

A little update for AC:

class DetectSquares:

    def __init__(self):
        self.x_to_points_dic = defaultdict(list)
        self.y_to_points_dic = defaultdict(list)
        self.points_cnt_dic = defaultdict(int)

    def add(self, point: List[int]) -> None:
        self.x_to_points_dic[point[0]].append(point)
        self.y_to_points_dic[point[1]].append(point)
        self.points_cnt_dic[tuple(point)] += 1
        
    def count(self, point: List[int]) -> int:
        res = 0
        x, y = point[0], point[1]
        for px in self.x_to_points_dic[x]:
            for py in self.y_to_points_dic[y]:
                if (py[0],px[1]) in self.points_cnt_dic and (py[0] != x) and (px[1] != y) and abs(py[0]-x) == abs(px[1]-y):
                    res += self.points_cnt_dic[(py[0],px[1])]
        return res

# Your DetectSquares object will be instantiated and called as such:
# obj = DetectSquares()
# obj.add(point)
# param_2 = obj.count(point)

x_to_points_dic and y_to_points_dic could be simplified with Counter():

class DetectSquares:

    def __init__(self):
        self.map = defaultdict(Counter)

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

    def count(self, point: List[int]) -> int:
        res = 0
        x, y = point

        if not y in self.map:
            return 0
        yCnt = self.map[y]

        for col, colCnt in self.map.items():
            if col != y:
                # 根据对称性，这里可以不用取绝对值
                d = col - y
                res += colCnt[x] * yCnt[x + d] * colCnt[x + d]
                res += colCnt[x] * yCnt[x - d] * colCnt[x - d]
        
        return res