Detect Squares

medium design hash map geometry

Problem

Design a class that works over a stream of lattice points with integer coordinates 0 ≤ x, y ≤ 1000. add(point) stores a point, and the same point may be added multiple times. count(point) returns how many axis-aligned squares of strictly positive area can be completed using the query point as one corner and three previously added points as the other three corners.

Every stored duplicate is a separate corner choice, so a corner that was added twice contributes twice as many squares.

Inputadd(3,10), add(11,2), add(3,2), count(11,10), count(14,8), add(11,2), count(11,10)

Output1, 0, 2

count(11,10) completes one square with (3,10), (11,2), (3,2); count(14,8) completes none; after a second add(11,2) the same square is counted twice.

operations (add x:y or count x:y, comma-separated)

from collections import defaultdict

class DetectSquares:
    def __init__(self):
        self.cnt = defaultdict(int)

    def add(self, point):
        self.cnt[tuple(point)] += 1

    def count(self, point):
        x, y = point
        total = 0
        for (px, py), c in list(self.cnt.items()):
            if px == x or abs(px - x) != abs(py - y):
                continue
            total += c * self.cnt[(x, py)] * self.cnt[(px, y)]
        return total

class DetectSquares {
  constructor() {
    this.cnt = new Map();
  }

  add(point) {
    const key = point[0] + "," + point[1];
    this.cnt.set(key, (this.cnt.get(key) || 0) + 1);
  }

  count(point) {
    const [x, y] = point;
    let total = 0;
    for (const [key, c] of this.cnt) {
      const [px, py] = key.split(",").map(Number);
      if (px === x || Math.abs(px - x) !== Math.abs(py - y)) continue;
      total += c * (this.cnt.get(x + "," + py) || 0) *
                   (this.cnt.get(px + "," + y) || 0);
    }
    return total;
  }
}

class DetectSquares {
    private Map<Integer, Integer> cnt = new HashMap<>();

    private int key(int x, int y) { return x * 1001 + y; }

    public void add(int[] point) {
        cnt.merge(key(point[0], point[1]), 1, Integer::sum);
    }

    public int count(int[] point) {
        int x = point[0], y = point[1], total = 0;
        for (Map.Entry<Integer, Integer> e : cnt.entrySet()) {
            int px = e.getKey() / 1001, py = e.getKey() % 1001;
            if (px == x || Math.abs(px - x) != Math.abs(py - y)) continue;
            total += e.getValue() * cnt.getOrDefault(key(x, py), 0)
                                  * cnt.getOrDefault(key(px, y), 0);
        }
        return total;
    }
}

class DetectSquares {
    unordered_map<int, int> cnt;
    int key(int x, int y) { return x * 1001 + y; }
public:
    void add(vector<int> point) {
        cnt[key(point[0], point[1])]++;
    }

    int count(vector<int> point) {
        int x = point[0], y = point[1], total = 0;
        for (auto& [k, c] : cnt) {
            int px = k / 1001, py = k % 1001;
            if (px == x || abs(px - x) != abs(py - y)) continue;
            if (cnt.count(key(x, py)) && cnt.count(key(px, y)))
                total += c * cnt[key(x, py)] * cnt[key(px, y)];
        }
        return total;
    }
};

Explanation

The whole structure is one hash map from coordinate to count. add just increments cnt[(x, y)], which transparently handles duplicates. All the thinking lives in count, and it rests on one geometric fact: an axis-aligned square is fully determined by one diagonal. If the query corner is (x, y) and some stored point (px, py) is the opposite corner, the other two corners are forced to be (x, py) and (px, y) — there is nothing left to choose.

So count(x, y) scans every distinct stored point as a diagonal candidate. A candidate is valid only when px ≠ x (positive area) and |px − x| = |py − y| (equal side lengths, i.e. the two points really sit on a 45° diagonal). For each valid candidate we look up the two forced corners in the map and add cnt[(px, py)] · cnt[(x, py)] · cnt[(px, y)] to the total.

The product is the counting trick: if the diagonal point was added twice and each forced corner once, there are 2 × 1 × 1 distinguishable squares, because any stored copy can serve as that corner. Missing corners contribute a factor of 0 and kill the term automatically.

Walking the default example: after add(3,10), add(11,2), add(3,2), the query count(11,10) scans the three stored points. Only (3, 2) passes the diagonal test (|3−11| = |2−10| = 8), and the forced corners (11, 2) and (3, 10) each have count 1, so the answer is 1 × 1 × 1 = 1. count(14,8) finds no candidate at all and returns 0. After a second add(11,2), the same diagonal now multiplies to 1 × 1 × 2 = 2.

Each count touches every distinct stored point once and does O(1) hash lookups per candidate, so it costs O(n) for n distinct points; add is O(1). With coordinates capped at 1000 there are at most 1001² distinct points, so the map — the only storage — is O(n).

Time: O(1) add, O(n) count Space: O(n)