Performant python datastructure for saving and querying coordinates

Question

for a map-like class I currently save some objects in a dictionary like:

{(3,5):<something>, (1,12):<whatever>, (17,2):<example>}

The keys are representing x/y coordinates in tuples. Not every coordinate in such an area is occupied, there are also "holes".

Some methods require to read out an area from that dictionary, such as "all objects from (2,5) to (6,10)". My present approach is to iterate over all keys and return all that are in the requested area. As these map-like dictionaries can get quite large I'm searching for a directer way for getting these areas. I thought about an OrdneredDict so I can stop the iteration if it gets past the requested area but maybe there is some better way?

Answer 1

The problem described is a typical problem that can be efficiently solved using a K-D tree.

A k-d tree (short for k-dimensional tree) is a space-partitioning data structure for organizing points in a k-dimensional space. In your case k is equal to 2. It is basically a binary search tree in which you alternate between the comparisons made at each level. For example if at a level you compare points on the basis of their x coordinate then on the next level you will compare points using the y coordinate.

So inserting a point in the tree can be done in O(logN). Insertion can be demonstrated by this image. Thus inserting N points in the data structure will take O(NlogN) time.

_{image taken from http://coursera.cs.princeton.edu/algs4/assignments/kdtree.html}

To find all points contained in a given query rectangle, start at the root and recursively search for points in both subtrees using the following pruning rule: if the query rectangle does not intersect the rectangle corresponding to a node, there is no need to explore that node (or its subtrees). A subtree is searched only if it might contain a point contained in the query rectangle.

Thus, the following running Java code solves this problem in the most optimum possible way. Each query for the points contained in a rectangle can be answered typically in O(R+logN) time. Where R is the number of points in the range and N is the number of points. However the worst case running time is O(R + sqrt(N))for pathological data.

    void run(){
        Scanner in = new Scanner(System.in);
        int numQueries = in.nextInt();
        Point2D root = null;
        for(int i=0; i<numQueries; i++){
            int type = in.nextInt();
            if(type == 0){   // Add a point
                double x = in.nextDouble();
                double y = in.nextDouble();
                root = addPoint(root, x, y, true);
            }else{
                double x1  = in.nextDouble();
                double y1 = in.nextDouble();
                double x2  = in.nextDouble();
                double y2 = in.nextDouble();
                System.out.println("the points in between the rectange with end points " + 
                     new Point2D(x1, y1) + " and " + new Point2D(x2, y2) + " are :");
                printPointsInRange(root, x1, y1, x2, y2, true);
            }
        }
    }


    // prints all points in the rectangle which has top left coordinates
    // (x1, y1) and bottom right coordinates (x2, y2)
    void printPointsInRange(Point2D root,
                        double x1, double y1, 
                        double x2, double y2, boolean vertical){
        if(root==null){
            return;
        }
        double x = root.x;
        double y = root.y;
        boolean outsideRange = Double.compare(x, x1)<0 || 
                                Double.compare(x, x2)>0 ||
                                 Double.compare(y, y1)>0 || 
                                 Double.compare(y, y2)<0;
        if(!outsideRange){
            System.out.println(root);
        }

        if(vertical){
            if(Double.compare(x, x1)<=0){ 
                // root lies left of left boundary or on the boundary
                printPointsInRange(root.right, x1, y1, x2, y2, !vertical);
            }else if(Double.compare(x, x2)>0){
                // root lies right of right boundary
                printPointsInRange(root.left, x1, y1, x2, y2, !vertical);
            }else if(Double.compare(x, x2)==0){
                // root lies on right boundary
                printPointsInRange(root.right, x1, y1, x2, y2, !vertical);
            }else{
                // root lies in between x1 and x2
                printPointsInRange(root.left, x1, y1, x2, y2, !vertical);
                printPointsInRange(root.right, x1, y1, x2, y2, !vertical);
            }
        }else{
            if(Double.compare(y, y2)<=0){ 
                // root lies below bottom boundary or on bottom boundary
                printPointsInRange(root.right, x1, y1, x2, y2, !vertical);
            }else if(Double.compare(y, y1)>0){
                // root lies above top boundary
                printPointsInRange(root.left, x1, y1, x2, y2, !vertical);
            }else if(Double.compare(y, y1)==0){
                // root lies on top boundary
                printPointsInRange(root.right, x1, y1, x2, y2, !vertical);
            }else{
                // root lies in between y1 and y2
                printPointsInRange(root.left, x1, y1, x2, y2, !vertical);
                printPointsInRange(root.right, x1, y1, x2, y2, !vertical);
            }
        }
    }

    Point2D addPoint(Point2D root, double x, double y, boolean vertical){
        if(root==null){
            return new Point2D(x, y);
        }
        if(vertical){   // vertical division
            double compare = Double.compare(root.x, x);
            if(compare<0){ // point is on left of root
                root.left = addPoint(root.left, x, y, !vertical);
            }else{        // point is on right of root or on root's x
                root.right = addPoint(root.right, x, y, !vertical);
            }
        }else{
            double compare = Double.compare(y, root.y);
            if(compare>0){    // point is above the root
                root.right = addPoint(root.right, x, y, !vertical);
            }else{            // point is below the root or on root's y 
                root.left = addPoint(root.left, x, y, !vertical);
            } 
        }
        return root;
    }

Answer 2

Instead of iterating over the entire dict, searching for those coordinates, you could generate all potential coordinates and check if they exist in the dict. This would be O(1) time vs your current solution of O(n) time.

Something like this would work (where grid is your dict),

def subset_coordinates(grid, top_left, bottom_right):
    a, b = top_left
    c, d = bottom_right
    for i in range(a, c+1):
        for j in range(b, d+1):
            value = grid.get((i, j))
            if value is not None:
                yield value

objects_in_subset = subset_coordinates(grid, (2, 5), (6, 10))