Detect Squares LeetCode solution in Python

1
# Runtime: 658 ms (Top 42.61%) | Memory: 16 MB (Top 39.48%)
2
"""
3
Store every points and number of their appearance
4
Group the points by their x value
5

6
Given a point(x1,y1) to count, try to find p2(x1,y2),p3(x2,y2),p4(x2,y1)
7
and add product of their appearances
8

9
Go through every points that has same x1(except the point that is the same as (x1,y1), the length of the side of the square is decided by abd(y2-y1).
10

11
Use the decided side length to calculate p3 and p4, see if they are in the dict. If do, add product of their appearances.
12

13
p3 p2 p3`
14

15
p4 p1 p4`
16

17
Notice that p3 and p4 can be on both left side and right side of side (p2,p1)
18

19
"""
20
from collections import defaultdict
21

22

23
class DetectSquares:
24

25
    def __init__(self):
26
        self.pts = defaultdict(int)
27
        self.by_x = defaultdict(set)
28

29
    def add(self, point: List[int]) -> None:
30
        self.pts[(point[0], point[1])] += 1
31
        self.by_x[point[0]].add((point[0], point[1]))
32

33
    def count(self, p1: List[int]) -> int:
34
        res = 0
35
        x1, y1 = p1[0], p1[1]
36
        if x1 in self.by_x:
37
            # p2:x1,y2
38
            for p2 in self.by_x[x1]:
39
                x2, y2 = p2
40
                if y1 == y2:
41
                    continue
42
                # length of side of square
43
                b = abs(y1 - y2)
44
                # left side
45
                p3 = (x1 - b, y2)
46
                p4 = (x1 - b, y1)
47
                if p3 in self.pts and p4 in self.pts:
48
                    res += self.pts[p2] * self.pts[p3] * self.pts[p4]
49
                # right side
50
                p3 = (x1 + b, y2)
51
                p4 = (x1 + b, y1)
52
                if p3 in self.pts and p4 in self.pts:
53
                    res += self.pts[p2] * self.pts[p3] * self.pts[p4]
54
        return res
55

56

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