c++algorithmgeometrycomputational-geometrykdtree

While finding count of points inside some shape using KD-Tree, do we need to check areas intersection, or just compare depth-appropriate properties?


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) {
    if (node == nullptr)
        return 0;

    return 1 + traverse(node->left) + traverse(node->right);
}

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);
        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);
        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


Solution

  • You certainly can, and in general you do.

    The normal way to query for all the points inside a given shape is to use the kD-tree as a broadphase for AABB testing, followed by a narrowphase test against the shape itself.

    So the first thing you'd do is identify the extents of the rectangle containing the shape. In the case of a circle obviously that's cx-cr < x < cx+cr and cy-cr < y < cy+cr. At an even (x-comparison) level there's no need to check the y extents, so there's your "depth-appropriate property". Of course, the node point itself needs the full comparison; luckily you only need to do it if both branches pass. So it's something like

    // even level
    bool leftPasses = node.x > xmin;
    bool rightPasses = node.x < xmax;
    
    if(leftPasses && rightPasses && narrowPhase(node.x, node.y)) count++;
    if(leftPasses) recurse(node.left);
    if(rightPasses) recurse(node.right);
    

    Note how I didn't do anything with the node boundaries. For something like this there's no reason to store them: If you're only doing a top-down traversal you can track them as you go. Also note the lack of a "check for containment and count the whole subtree" strategy, which is not useful from an asymptotic performance standpoint.