Need help optimizing orthogonal range search for a static kd tree

Question

I'm trying to optimise my implemention of a static kd tree to perform orthogonal range searches in c++. My problem: The code performs slowly even for a small number of queries when the number of points is around 10^5.

I've constructed the tree based on http://en.wikipedia.org/wiki/K-d_tree (i.e, It's a static kd tree where the data is stored both in the nodes and the leaves)

And the orthogonal range searching algorithm based on Ch 11 of http://www.cs.umd.edu/~mount/754/Lects/754lects.pdf

Relevant algorithm: (My implementation returns the points instead of count)

int rangeCount(Range Q, KDNode t, Rectangle C)
(1) if (t is a leaf)
  (a) if (Q contains t) return 1,
  (b) else return 0.
(2) if (t is not a leaf)
  (a) if (C does not intersect Q) return 0.
  (b) else if (C is a subset of Q) return t:size.
  (c) else, split C along t’s cutting dimension and cutting value, letting C1 and C2 be the two rectangles. Return (rangeCount(Q; t:left; C1) + rangeCount(Q; t:right; C2)).

My code (The relevant function is the 2nd recursive_query())

#include <iostream>
#include <algorithm>
#include <cstdio>
#include <vector>
#include <random>
#include <chrono>
using namespace std;
typedef long long ll;

typedef pair<int, int> T;
typedef pair<pair<int, int>, pair<int, int> > R;

bool comp_x(const T &a, const T &b)
{
    return a.first < b.first;
}
bool comp_y(const T &a, const T &b)
{
    return a.second < b.second;
}

template<int k=2>
class kd_tree
{
private:
    vector<T> array; // array[0] is a dummy
    int N;

    void recursive_build(int depth, int node, int begin, int end, vector<T> &elements)
    {
        if(end < begin)
            return;
        if(begin == end){
            array[node] = elements[begin];
            return;
        }

        int axis = depth % k;
        int median;

        // Find median for axis (Todo: Implement quickselect)
        if(axis == 0) // x-axis
        {
            sort(elements.begin()+begin, elements.begin()+end+1, comp_x);
            median = (begin+end+1)/2;
        }
        else
        {
            sort(elements.begin()+begin, elements.begin()+end+1, comp_y);
            median = (begin+end+1)/2;
        }

        array[node] = elements[median];
        recursive_build(depth+1, 2*node, begin, median-1, elements);
        recursive_build(depth+1, 2*node+1, median+1, end, elements);
    }
    void return_subtree(int node, vector<T> &query_list)
    {
        if(node > N || array[node].first == 0)
            return;

        query_list.push_back(array[node]);
        return_subtree(2*node, query_list);
        return_subtree(2*node+1, query_list);
    }
    void recursive_query(int depth, int node, int x1, int y1, int x2, int y2, vector<T> &query)
    {
        if(node > N)
            return;

        int axis = depth % k;

        if(axis == 0) // x-axis
        {
            if(array[node].first < x1)
            {
                recursive_query(depth+1, 2*node+1, x1, y1, x2, y2, query);
            }
            else if(array[node].first > x2)
            {
                recursive_query(depth+1, 2*node, x1, y1, x2, y2, query);
            }
            else
            {
                if(array[node].second >= y1 && array[node].second <= y2)
                    query.push_back(array[node]);
                recursive_query(depth+1, 2*node, x1, y1, x2, y2, query);
                recursive_query(depth+1, 2*node+1, x1, y1, x2, y2, query);
            }
        }
        else
        {
            if(array[node].second < y1)
            {
                recursive_query(depth+1, 2*node+1, x1, y1, x2, y2, query);
            }
            else if(array[node].second > y2)
            {
                recursive_query(depth+1, 2*node, x1, y1, x2, y2, query);
            }
            else
            {
                if(array[node].first >= x1 && array[node].first <= x2)
                    query.push_back(array[node]);
                recursive_query(depth+1, 2*node, x1, y1, x2, y2, query);
                recursive_query(depth+1, 2*node+1, x1, y1, x2, y2, query);
            }
        }

    }
    void recursive_query(int depth, int node, R cell, R query, vector<T> &query_list)
    {
        int left = 2*node;
        int right = 2*node+1;

        // node is a leaf node
        if(array[left].first == 0 && array[right].first == 0) 
        {
            // if the leaf lies inside the query rectangle then add it to query_list
            if(array[node].first >= query.first.first && array[node].first <= query.second.first
                && array[node].second >= query.first.second && array[node].second <= query.second.second)
            {
                query_list.push_back(array[node]);
            }
            return;
        }

        // node is not a leaf

        // cell doesnt intersect the query
        if(cell.first.first > query.second.first || cell.second.first < query.first.first
            || cell.first.second  > query.second.second || cell.second.second < query.first.second)
        {
            return;
        }

        // cell is a subset of query
        if(cell.first.first >= query.first.first && cell.second.first <= query.second.first
            && cell.first.second >= query.first.second && cell.second.second <= query.second.second)
        {
            return_subtree(node, query_list);
            return;
        }

        // if the node lies within bounds then add it to query_list
        if(array[node].first >= query.first.first && array[node].first <= query.second.first
            && array[node].second >= query.first.second && array[node].second <= query.second.second)
            {
                query_list.push_back(array[node]);
            }

        depth = depth % k;
        if(depth == 0) // splitting planes is the x-axis
        {
            if(array[left].first != 0){
                R cell1 = make_pair(cell.first, make_pair(array[node].first, cell.second.second));
                recursive_query(depth+1, left, cell1, query, query_list);
            }
            if(array[right].first != 0){
                R cell2 = make_pair(make_pair(array[node].first, cell.first.second), cell.second);
                recursive_query(depth+1, right, cell2, query, query_list);
            }
        }
        else // splitting plane is the y-axis
        {
            if(array[left].first != 0){
                R cell1 = make_pair(cell.first, make_pair(cell.second.first, array[node].second));
                recursive_query(depth+1, left, cell1, query, query_list);
            }
            if(array[right].first != 0){
                R cell2 = make_pair(make_pair(cell.first.first, array[node].second), cell.second);
                recursive_query(depth+1, right, cell2, query, query_list);
            }
        }
    }
public:
    kd_tree(vector<T> &elements)
    {
        N = elements.size();
        array.resize(2*k*N+1);  
        recursive_build(0, 1, 0, N-1, elements);
    }
    void orthogonal_query(int x1, int y1, int x2, int y2, vector<T> &query_list)
    {
//      recursive_query(0, 1, x1, y1, x2, y2, query_list);
        recursive_query(0, 1, make_pair(make_pair(1, 1), make_pair(300000, 300000)), make_pair(make_pair(x1, y1), make_pair(x2, y2)), query_list);
    }

};


int main(void)
{
    std::mt19937_64 generator;
    std::uniform_int_distribution<int> distribution(0, 300000);

    int n, q, x, y, d, count;
//  scanf("%d %d", &n, &q);
    n = 30000;
    q = 2000;

    vector<pair<int, int> > points, query;
    vector<pair<int, int> >::iterator itr;

    // Input Phase
    while (n--){
//      scanf("%d %d", &x, &y);
        x = distribution(generator);
        y = distribution(generator);
        points.push_back(make_pair(x, y));
    }

    kd_tree<2> tree(points);

    std::chrono::time_point<std::chrono::high_resolution_clock> start, end;
    start = std::chrono::high_resolution_clock::now();

    // Query Loop
    while(q--){

        query.clear();
//      scanf("%d %d %d", &x, &y, &d);
        x = distribution(generator);
        y = distribution(generator);
        d = distribution(generator);

        int d1 = x+y+d;

        tree.orthogonal_query(x, y, x+d, y+d, query);

//      printf("%d\n", count);
    }

    end = std::chrono::high_resolution_clock::now();
    ll elapsed_time = std::chrono::duration_cast<std::chrono::milliseconds>(end-start).count();
    cout << "\nElapsed Time: " << elapsed_time << "ms\n";

    return 0;
}

Answer 1

You could make the return_subtree function non-recursive by using a stack object by yourself:

void return_subtree(int node, vector<T> &query_list)
{
    if (node > N) { return; }

    std::stack<int> st;
    st.push(node);

    while (not st.empty()) {
        int top = st.top();
        st.pop();

        // if array[top].first == 0, then the recursion ends...
        // ... BUT it still need be returned!
        if (top > N) { continue; }

        query_list.push_back(array[top]);
        if(array[top].first == 0) { continue; }

        st.push(2 * top + 1);
        st.push(2*node+1);
    }
}

This could lead to a speed up because return_subtree cannot be tail-call optimized (or at least, only the latest call may be optimized), and therefore the compiler is probably unable to come up with this optimization by itself...

... which of course brings us to the crux of the matter: turn your compiler optimizations on (-Os or -O2 flag on the command line for gcc/clang).

Other than that, I don't see any glaring inefficiency.