Let’s say we have to find count of points inside some shape in 2D plane.

We can use KD-Tree to solve this. The KD-Tree will partition 2D plane into rectangles, and then we can check intersections between some node boundaries (which looks like rectangle) and the shape inside which we want to count points.

That seems logical. I have written code for solving problem like this, where my shape is circle.

static bool is_point_inside_circle(int x, int y, int cx, int cy, int cr) {
    return (x - cx) * (x - cx) + (y - cy) * (y - cy) <= cr * cr;
}

struct RectBoundary {
    int top;
    int right;
    int bottom;
    int left;

    RectBoundary() : top(500), right(500), bottom(-500), left(-500) { }
    RectBoundary(int t, int r, int b, int l) : top(t), right(r), bottom(b), left(l) { }

    bool is_inside_circle(int cx, int cy, int cr) const {
        return is_point_inside_circle(left, top, cx, cy, cr)
            && is_point_inside_circle(right, top, cx, cy, cr)
            && is_point_inside_circle(right, bottom, cx, cy, cr)
            && is_point_inside_circle(left, bottom, cx, cy, cr);
    }

    bool intersects_with_circle(int cx, int cy, int cr) const {
        return cx - cr <= left
            || cx + cr >= right
            || cy - cr <= bottom
            || cy + cr >= top
            || (cx >= left && cx <= right && cy <= top && cy >= bottom);
    }
};

struct KDNode {
    int x, y;
    int depth;
    RectBoundary boundary;

    KDNode* left = nullptr;
    KDNode* right = nullptr;

    bool is_leaf() {
        return left == nullptr && right == nullptr;
    }
};

KDNode* build(vector<vector<int>>& points, int l, int r, int depth) {
    if (l > r)
        return nullptr;

    if (depth % 2 == 0) {
        sort(points.begin() + l, points.begin() + r + 1);
    }
    else {
        sort(points.begin() + l, points.begin() + r + 1, [](const vector<int>& p1, const vector<int>& p2)
            {
                return p1[1] < p2[1];
            });
    }

    const int mid = (l + r) / 2;

    KDNode* node = new KDNode();
    node->depth = depth;
    node->x = points[mid][0];
    node->y = points[mid][1];

    node->left = build(points, l, mid - 1, depth + 1);
    node->right = build(points, mid + 1, r, depth + 1);

    return node;
}

void set_boundaries(KDNode* node, int depth) {
    if (node == nullptr)
        return;

    if (depth % 2 == 0) {
        if (node->left) {
            node->left->boundary = node->boundary;
            node->left->boundary.right = node->x;
        }

        if (node->right) {
            node->right->boundary = node->boundary;
            node->right->boundary.left = node->x;
        }
    }
    else {
        if (node->left) {
            node->left->boundary = node->boundary;
            node->left->boundary.top = node->y;
        }

        if (node->right) {
            node->right->boundary = node->boundary;
            node->right->boundary.bottom = node->y;
        }
    }

    set_boundaries(node->left, depth + 1);
    set_boundaries(node->right, depth + 1);
}

int traverse(KDNode* node, int cx, int cy, int cr) {
    if (node == nullptr)
        return 0;

    return 1 + traverse(node->left, cx, cy, cr) + traverse(node->right, cx, cy, cr);
}

int count_points_inside_circle(KDNode* node, int cx, int cy, int cr) {
    if (node->is_leaf()) {
        return is_point_inside_circle(node->x, node->y, cx, cy, cr);
    }

    int count = is_point_inside_circle(node->x, node->y, cx, cy, cr);

    if (node->left) {
        if (node->boundary.is_inside_circle(cx, cy, cr))
            count += traverse(node->left, cx, cy, cr);
        else if (node->boundary.intersects_with_circle(cx, cy, cr))
            count += count_points_inside_circle(node->left, cx, cy, cr);
    }

    if (node->right) {
        if (node->boundary.is_inside_circle(cx, cy, cr))
            count += traverse(node->right, cx, cy, cr);
        else if (node->boundary.intersects_with_circle(cx, cy, cr))
            count += count_points_inside_circle(node->right, cx, cy, cr);
    }

    return count;
}

vector<int> countPoints(vector<vector<int>>& points, vector<vector<int>>& queries) {
    KDNode* root = build(points, 0, points.size() - 1, 0);
    set_boundaries(root, 0);

    vector<int> result(queries.size());
    for (int i = 0; i < queries.size(); ++i)
        result[i] = count_points_inside_circle(root, queries[i][0], queries[i][1], queries[i][2]);

    delete root;
    return result;
}

After writing this, I realized that my intersects_with_circle is most probably wrong, as it assumes fixed relative position of rectangle and circle. I still can’t figure out appropriate formula for checking intersection between circle and rectangle.

But what came to my mind, is: do we really need to compare whole area of shapes?

I mean, given rectangle and any shape, what if we compare only properties of shape that are appropriate for given KD-Tree depth? Lets say we have 2D plane, and current depth is 6. It’s even, so we will use x property. Can we just get the middle point of the shape inside which we want to count points, check x orientation against this middle point using cross product, and then get the most left or most right point of shape (depends on orientation), and check if this “most point” x property is lesser or greater (depends on orientation also) than x of our rectangle?

Could it work? Can we like isolate to one property at the moment, instead of comparing whole areas? I can imagine that the given shape “most point” x property can be not valid without checking y also, but then we will eventually check also the y property and optionally reject some specific area (one additional step, from which we could have got rid of earlier), as the properties alternate based on KD-Tree depth.

I’m not sure if I understand KD-Tree properly and it’s possible that this question doesn’t make any sense. Sorry for that