Search in sources :

Example 1 with KDTreeNode

use of com.airbnb.aerosolve.core.KDTreeNode in project aerosolve by airbnb.

the class KDTreeModel method queryBox.

// Returns the indices of all node overlapping the box
public ArrayList<Integer> queryBox(double minX, double minY, double maxX, double maxY) {
    ArrayList<Integer> idx = new ArrayList<>();
    if (nodes == null)
        return idx;
    Stack<Integer> stack = new Stack<Integer>();
    stack.push(0);
    while (!stack.isEmpty()) {
        int currIdx = stack.pop();
        idx.add(currIdx);
        KDTreeNode node = nodes[currIdx];
        switch(node.nodeType) {
            case X_SPLIT:
                {
                    if (minX < node.splitValue) {
                        stack.push(node.leftChild);
                    }
                    if (maxX >= node.splitValue) {
                        stack.push(node.rightChild);
                    }
                }
                break;
            case Y_SPLIT:
                {
                    if (minY < node.splitValue) {
                        stack.push(node.leftChild);
                    }
                    if (maxY >= node.splitValue) {
                        stack.push(node.rightChild);
                    }
                }
            case LEAF:
                break;
        }
    }
    return idx;
}
Also used : KDTreeNode(com.airbnb.aerosolve.core.KDTreeNode) ArrayList(java.util.ArrayList) Stack(java.util.Stack)

Example 2 with KDTreeNode

use of com.airbnb.aerosolve.core.KDTreeNode in project aerosolve by airbnb.

the class KDTreeModelTest method getTestNodes.

//                    4
//         |--------------- y = 2
//  1      | 2       3
//     x = 1
public static KDTreeNode[] getTestNodes() {
    KDTreeNode parent = new KDTreeNode();
    parent.setNodeType(KDTreeNodeType.X_SPLIT);
    parent.setSplitValue(1.0);
    parent.setLeftChild(1);
    parent.setRightChild(2);
    KDTreeNode one = new KDTreeNode();
    one.setNodeType(KDTreeNodeType.LEAF);
    KDTreeNode two = new KDTreeNode();
    two.setNodeType(KDTreeNodeType.Y_SPLIT);
    two.setSplitValue(2.0);
    two.setLeftChild(3);
    two.setRightChild(4);
    KDTreeNode three = new KDTreeNode();
    three.setNodeType(KDTreeNodeType.LEAF);
    KDTreeNode four = new KDTreeNode();
    four.setNodeType(KDTreeNodeType.LEAF);
    KDTreeNode[] arr = { parent, one, two, three, four };
    return arr;
}
Also used : KDTreeNode(com.airbnb.aerosolve.core.KDTreeNode)

Example 3 with KDTreeNode

use of com.airbnb.aerosolve.core.KDTreeNode in project aerosolve by airbnb.

the class KDTreeModel method next.

private int next(int currIdx, double x, double y) {
    KDTreeNode node = nodes[currIdx];
    int nextIndex = -1;
    switch(node.nodeType) {
        case X_SPLIT:
            {
                if (x < node.splitValue) {
                    nextIndex = node.leftChild;
                } else {
                    nextIndex = node.rightChild;
                }
            }
            break;
        case Y_SPLIT:
            {
                if (y < node.splitValue) {
                    nextIndex = node.leftChild;
                } else {
                    nextIndex = node.rightChild;
                }
            }
            break;
        default:
            assert (node.nodeType == LEAF);
            break;
    }
    return nextIndex;
}
Also used : KDTreeNode(com.airbnb.aerosolve.core.KDTreeNode)

Aggregations

KDTreeNode (com.airbnb.aerosolve.core.KDTreeNode)3 ArrayList (java.util.ArrayList)1 Stack (java.util.Stack)1