diff --git a/orx-kdtree/src/demo/kotlin/DemoKNearestNeighbour01.kt b/orx-kdtree/src/demo/kotlin/DemoKNearestNeighbour01.kt new file mode 100644 index 00000000..0204a12b --- /dev/null +++ b/orx-kdtree/src/demo/kotlin/DemoKNearestNeighbour01.kt @@ -0,0 +1,35 @@ +import org.openrndr.application +import org.openrndr.color.ColorRGBa +import org.openrndr.extra.kdtree.buildKDTree +import org.openrndr.extra.kdtree.findKNearest +import org.openrndr.extra.kdtree.vector2Mapper +import org.openrndr.math.Vector2 +import org.openrndr.shape.LineSegment + +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) + + extend { + drawer.circles(points, 5.0) + + val kNearest = findKNearest(tree, mouse.position, k=7, dimensions = 2, ::vector2Mapper) + drawer.fill = ColorRGBa.RED + drawer.stroke = ColorRGBa.RED + drawer.strokeWeight = 2.0 + drawer.circles(kNearest, 7.0) + drawer.lineSegments(kNearest.map { LineSegment(mouse.position, it) }) + } + } + } +} \ 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 e39f3a8e..8b0ff905 100644 --- a/orx-kdtree/src/main/kotlin/KDTree.kt +++ b/orx-kdtree/src/main/kotlin/KDTree.kt @@ -165,7 +165,6 @@ private fun sqrDistance(left: T, right: T, dimensions: Int, mapper: (T, Int) fun findAllNodes(root: KDTreeNode): List> { val stack = Stack>() val all = ArrayList>() - stack.empty() stack.push(root) while (!stack.isEmpty()) { val node = stack.pop() @@ -184,6 +183,54 @@ fun findAllNodes(root: KDTreeNode): List> { } +fun findKNearest( + root: KDTreeNode, + item: T, + k: Int, + dimensions: Int, + mapper: (T, Int) -> Double +): List { + // max-heap with size k + val queue = PriorityQueue, Double>>(k + 1) { + nodeA, nodeB -> compareValues(nodeB.second, nodeA.second) + } + + fun nearest(node: KDTreeNode?, item: T) { + if (node != null) { + val dimensionValue = mapper(item, node.dimension) + val route: Int = if (dimensionValue < node.median) { + nearest(node.children[0], item) + 0 + } else { + nearest(node.children[1], item) + 1 + } + + val distance = sqrDistance(item, node.item + ?: throw IllegalStateException("item is null"), dimensions, mapper) + + if (queue.size < k || distance < queue.peek().second) { + queue.add(Pair(node, distance)) + if (queue.size > k) { + queue.poll() + } + } + + val d = abs(node.median - dimensionValue) + if (d * d < queue.peek().second || queue.size < k) { + nearest(node.children[1 - route], item) + } + } + } + + nearest(root, item) + + return generateSequence { queue.poll() } + .map { it.first.item } + .filterNotNull() + .toList().reversed() +} + fun findNearest(root: KDTreeNode, item: T, dimensions: Int, mapper: (T, Int) -> Double): T? { var nearest = java.lang.Double.POSITIVE_INFINITY var nearestArg: KDTreeNode? = null