diff --git a/orx-kdtree/src/demo/kotlin/DemoRangeQuery01.kt b/orx-kdtree/src/demo/kotlin/DemoRangeQuery01.kt new file mode 100644 index 00000000..3a86b965 --- /dev/null +++ b/orx-kdtree/src/demo/kotlin/DemoRangeQuery01.kt @@ -0,0 +1,39 @@ +import org.openrndr.application +import org.openrndr.color.ColorRGBa +import org.openrndr.extra.kdtree.buildKDTree +import org.openrndr.extra.kdtree.findAllInRange +import org.openrndr.extra.kdtree.vector2Mapper +import org.openrndr.math.Vector2 + + +fun main() { + application { + + configure { + width = 1080 + height = 720 + } + + program { + val points = MutableList(1000) { + Vector2(Math.random() * width, Math.random() * height) + } + val tree = buildKDTree(points, 2, ::vector2Mapper) + val radius = 50.0 + + extend { + drawer.circles(points, 5.0) + + val allInRange = findAllInRange(tree, mouse.position, maxDistance = radius, dimensions = 2, ::vector2Mapper) + drawer.fill = ColorRGBa.PINK + drawer.stroke = ColorRGBa.PINK + drawer.strokeWeight = 2.0 + drawer.circles(allInRange, 7.0) + + drawer.fill = null + drawer.strokeWeight = 1.0 + drawer.circle(mouse.position, radius) + } + } + } +} \ No newline at end of file diff --git a/orx-kdtree/src/main/kotlin/KDTree.kt b/orx-kdtree/src/main/kotlin/KDTree.kt index 8b0ff905..5a7705ee 100644 --- a/orx-kdtree/src/main/kotlin/KDTree.kt +++ b/orx-kdtree/src/main/kotlin/KDTree.kt @@ -268,6 +268,47 @@ fun findNearest(root: KDTreeNode, item: T, dimensions: Int, mapper: (T, I return nearestArg?.item } +fun findAllInRange( + root: KDTreeNode, + item: T, + maxDistance: Double, + dimensions: Int, + mapper: (T, Int) -> Double +) : List { + + val sqrMaxDist = maxDistance * maxDistance + val queue = kotlin.collections.ArrayDeque?>() + queue.add(root) + val results = mutableListOf() + + while (queue.isNotEmpty()) { + val node = queue.removeFirst() + if (node != null) { + val dimensionValue = mapper(item, node.dimension) + val distance = sqrDistance(item, node.item + ?: throw IllegalStateException("item is null"), dimensions, mapper) + if (distance <= sqrMaxDist) { + results.add(node.item) + } + + val route: Int = if (dimensionValue < node.median) { + queue.add(node.children[0]) + 0 + } else { + queue.add(node.children[1]) + 1 + } + + val d = abs(node.median - dimensionValue) + if (d * d <= sqrMaxDist) { + queue.add(node.children[1 - route]) + } + } + } + + return results.filterNotNull() +} + fun insert(root: KDTreeNode, item: T, dimensions: Int, mapper: (T, Int) -> Double): KDTreeNode { val stack = Stack>() stack.push(root)