Implement Quadtree structure with geometric shapes

Library/Geometry/quadtree.cpp

@@ -0,0 +1,153 @@
+#include <bits/stdc++.h> //LATEX_IGNORED_LINE
+using namespace std;
+
+const int MXN = 1e5+100;
+
+struct Point {
+    int x, y;
+
+    bool const operator<(const Point& o) const {
+        return make_tuple(x, y) < make_tuple(o.x, o.y);
+    }
+};
+
+struct AABB {
+    int x1, y1, x2, y2;
+
+    AABB() : x1(0), y1(0), x2(MXN), y2(MXN) {} 
+    AABB(int x1, int y1, int x2, int y2) : x1(x1), y1(y1), x2(x2), y2(y2) {}
+
+    bool intersects(const AABB& other) const {
+        bool overlap_x = (x1 <= other.x2) && (x2 >= other.x1);
+        bool overlap_y = (y1 <= other.y2) && (y2 >= other.y1);
+        return overlap_x && overlap_y;
+    }
+
+    bool contains(int x, int y) const {
+        return (x >= x1 && x <= x2 && y >= y1 && y <= y2);
+    }
+
+    array<AABB, 4> subdivide() const {
+        int x_center = x1 + (x2 - x1) / 2;
+        int y_center = y1 + (y2 - y1) / 2;
+        return {
+            AABB{x1, y1, x_center, y_center},     // NW
+            AABB{x_center, y1, x2, y_center},     // NE
+            AABB{x1, y_center, x_center, y2},     // SW
+            AABB{x_center, y_center, x2, y2}      // SE
+        };
+    }
+};
+
+
+struct Circle {
+    Point c;
+    int r, w;
+    AABB box;
+
+    Circle(Point c, int r, int w) 
+        : c(c), r(r), w(w), box(c.x - r, c.y - r, c.x + r, c.y + r) {}
+
+    bool const operator<(const Circle& o) const {
+        return make_tuple(c.x, c.y, r) < make_tuple(o.c.x, o.c.y, o.r);
+    }
+};
+
+struct Square {
+    Point p1, p2;
+    int w;
+    AABB box;
+
+    Square(Point p1, Point p2)
+        : p1(p1), p2(p2), box(min(p1.x, p2.x), min(p1.y, p2.y), 
+                              max(p1.x, p2.x), max(p1.y, p2.y)) {}
+
+    bool const operator<(const Square& o) const {
+        return make_tuple(p1, p2) < make_tuple(o.p1, o.p2);
+    }
+};
+
+struct Triangle {
+    Point p1, p2, p3;
+    int w;
+    AABB box;
+
+    Triangle(Point p1, Point p2, Point p3)
+        : p1(p1), p2(p2), p3(p3),
+          box(min({p1.x, p2.x, p3.x}), min({p1.y, p2.y, p3.y}), max({p1.x, p2.x, p3.x}), max({p1.y, p2.y, p3.y})) {}
+
+    bool const operator<(const Triangle& o) const {
+        return make_tuple(p1, p2, p3) < make_tuple(o.p1, o.p2, o.p3);
+    }
+};
+
+template <typename T>
+class Quadtree {
+  public:
+    static const int MIN_OBJECTS = 8;
+    static const int MAX_DEPTH = 16;
+
+    Quadtree(const AABB& root_bounds, const vector<T>& all_objects) 
+        : bounds(root_bounds), depth(0) {
+
+        vector<const T*> object_ptrs;
+        object_ptrs.reserve(all_objects.size());
+        for (const auto& obj : all_objects) {
+            object_ptrs.push_back(&obj);
+        }
+
+        build(object_ptrs);
+    }
+
+    const vector<const T*>& query(int x, int y) const {
+        const Quadtree* node = this;
+
+        while (node->is_branch) {
+            int child_index = node->get_child_index(x, y);
+            if (child_index == -1) break;
+            node = node->children[child_index].get();
+        }
+
+        return node->objects;
+    }
+
+  private:
+    AABB bounds;
+    int depth;
+    bool is_branch = false;
+    int x_center, y_center;
+
+    array<unique_ptr<Quadtree<T>>, 4> children;
+
+    vector<const T*> objects;
+
+    Quadtree(const AABB& box, int d) : bounds(box), depth(d) {}
+
+    void build(const vector<const T*>& parent_objects) {
+        x_center = bounds.x1 + (bounds.x2 - bounds.x1) / 2;
+        y_center = bounds.y1 + (bounds.y2 - bounds.y1) / 2;
+
+        for (const T* obj : parent_objects) {
+            if (bounds.intersects(obj->box)) {
+                objects.push_back(obj);
+            }
+        }
+
+        if (objects.size() > MIN_OBJECTS && depth < MAX_DEPTH) {
+            is_branch = true;
+            auto child_bounds = bounds.subdivide();
+
+            for (int i = 0; i < 4; ++i) {
+                children[i] = unique_ptr<Quadtree<T>>(new Quadtree<T>(child_bounds[i], depth + 1));
+                children[i]->build(objects);
+            }
+
+            objects.clear();
+        }
+    }
+
+    int get_child_index(int x, int y) const {
+        if (y <= y_center) return (x <= x_center) ? 0 : 1; // 0 (NW) ou 1 (NE)
+        else return (x <= x_center) ? 2 : 3; // 2 (SW) ou 3 (SE)   
+    }
+};