Detect Squares

NeetCode: Detect Squares · LeetCode: 2013. Detect Squares · Difficulty: Medium

Design a data structure that supports:

add(int[] point): add a point [x, y] (duplicates are allowed and count separately).
count(int[] point): count the ways to choose three stored points that, together with the query point, form an axis-aligned square with positive area.

Example

Text

add [3, 10], add [11, 2], add [3, 2]
count [11, 10]  -> 1
count [14, 8]   -> 0
add [11, 2]                      (a duplicate)
count [11, 10]  -> 2

Brute force: scan all stored points

Keep every added point in a list. For count(px, py), a stored point (qx, qy) can be the diagonal opposite corner only if |qx - px| == |qy - py| and the distance is non-zero. Each such candidate fixes the other two corners at (px, qy) and (qx, py); count how many times each is stored by scanning the list again. The sum over all candidates is the answer.

Java

import java.util.ArrayList;
import java.util.List;

class DetectSquares {
    private final List<int[]> points = new ArrayList<>();

    public DetectSquares() {
    }

    public void add(int[] point) {
        points.add(new int[] {point[0], point[1]});
    }

    private int occurrences(int x, int y) {
        int c = 0;
        for (int[] q : points) {
            if (q[0] == x && q[1] == y) c++;
        }
        return c;
    }

    public int count(int[] point) {
        int px = point[0], py = point[1];
        int total = 0;
        for (int[] q : points) {
            int dx = Math.abs(q[0] - px);
            int dy = Math.abs(q[1] - py);
            if (dx == 0 || dx != dy) continue;
            total += occurrences(px, q[1]) * occurrences(q[0], py);
        }
        return total;
    }
}

public class Main {
    public static void main(String[] args) {
        DetectSquares ds = new DetectSquares();
        ds.add(new int[] {3, 10});
        ds.add(new int[] {11, 2});
        ds.add(new int[] {3, 2});
        System.out.println(ds.count(new int[] {11, 10}));
        System.out.println(ds.count(new int[] {14, 8}));
        ds.add(new int[] {11, 2});
        System.out.println(ds.count(new int[] {11, 10}));
    }
}

Output

1
0
2

Time: add is O(1); count is O(n²) for n stored points. Space: O(n).

Optimal: hash map of point counts

Store a HashMap from point to count, so a duplicate is one entry with a larger count. An int[] cannot be the key because arrays compare by identity, so the key is a record Point(int x, int y). A record is a short class whose equals and hashCode compare its fields, which makes it a valid map key.

In count, iterate over the distinct stored points as the diagonal candidate; the two remaining corners are O(1) map lookups. A candidate with multiplicity c contributes c * count(px, qy) * count(qx, py). The inner scans disappear, so count is linear in the number of distinct points. HashMap iteration order is unspecified, but the answer is a sum, so the order does not change it.

Java

import java.util.HashMap;
import java.util.Map;

class DetectSquares {
    private record Point(int x, int y) {
    }

    private final Map<Point, Integer> counts = new HashMap<>();

    public DetectSquares() {
    }

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

    public int count(int[] point) {
        int px = point[0], py = point[1];
        int total = 0;
        for (Map.Entry<Point, Integer> e : counts.entrySet()) {
            int qx = e.getKey().x(), qy = e.getKey().y();
            int dx = Math.abs(qx - px);
            int dy = Math.abs(qy - py);
            if (dx == 0 || dx != dy) continue;
            int c1 = counts.getOrDefault(new Point(px, qy), 0);
            int c2 = counts.getOrDefault(new Point(qx, py), 0);
            total += e.getValue() * c1 * c2;
        }
        return total;
    }
}

public class Main {
    public static void main(String[] args) {
        DetectSquares ds = new DetectSquares();
        ds.add(new int[] {3, 10});
        ds.add(new int[] {11, 2});
        ds.add(new int[] {3, 2});
        System.out.println(ds.count(new int[] {11, 10}));
        System.out.println(ds.count(new int[] {14, 8}));
        ds.add(new int[] {11, 2});
        System.out.println(ds.count(new int[] {11, 10}));
    }
}

Output

1
0
2

Time: add is O(1); count is O(k) for k distinct points. Space: O(k).