KdPlaceTree.kt

class KdPlaceTree(points : List<XYZPoint>) : KdTree<KdTree.XYZPoint>(points) {

    var counter : Int = 0

    fun radialSearch(distance: Double, query : XYZPoint) : List<XYZPoint> {
        if (root == null)
            return emptyList()

        val results = TreeSet<KdNode>(EuclideanComparator(query))

        var prev : KdNode? = null
        var node : KdNode? = root

        // find the closest leaf node
        while(node != null) {
            if (KdNode.compareTo(node.depth, node.k, query, node.point) <= 0) {
                // lesser
                prev = node
                node = node.lesser
            } else {
                // greater
                prev = node
                node = node.greater
            }
        }

        val leaf = prev

        if (leaf != null) {
            // if the closest node is too far, return empty result right away.
            if (leaf.point.euclideanDistance(query) > distance)
                return emptyList()
            else {

                val examined = hashSetOf<KdNode>()

                node = leaf

                while (node != null) {
                    searchPointsWithinCertainDistance(
                            distance = distance,
                            query = query,
                            examined = examined,
                            results = results,
                            node = node
                    )
                    node = node.parent
                }

                return results.map { result -> result.point }
            }

        } else {
            return emptyList()
        }

    }

    private fun searchPointsWithinCertainDistance(
            distance : Double,
            query : XYZPoint,
            node : KdNode,
            results : TreeSet<KdNode>,
            examined : HashSet<KdNode>
    ) {

        counter++

        examined.add(node)

        // calculating distance between this node's query and target query.
        val nodeDistance = node.point.euclideanDistance(query)

        if (nodeDistance <= distance)
            results.add(node)

        val lastNode = results.last()

        val lastDistance = lastNode.point.euclideanDistance(query)

        val axis = node.depth % node.k

        val lesser = node.lesser
        val greater = node.greater

        if (lesser != null && !examined.contains(lesser)) {

            examined.add(lesser)

            val nodePoint : Double
            val pointPlusDistance : Double

            when(axis) {
                X_AXIS -> {
                    nodePoint = node.point.x
                    pointPlusDistance = query.x - lastDistance
                }
                Y_AXIS -> {
                    nodePoint = node.point.y
                    pointPlusDistance = query.y - lastDistance
                }
                else -> {
                    nodePoint = node.point.z
                    pointPlusDistance = query.z - lastDistance
                }
            }

            val lineIntersectsCube = pointPlusDistance <= nodePoint

            if (lineIntersectsCube)
                searchPointsWithinCertainDistance(
                        distance = distance,
                        query = query,
                        node = lesser,
                        results = results,
                        examined = examined
                )
        }

        if (greater != null && !examined.contains(greater)) {

            examined.add(greater)

            val nodePoint : Double
            val pointPlusDistance : Double

            when(axis) {
                X_AXIS -> {
                    nodePoint = node.point.x
                    pointPlusDistance = query.x + lastDistance
                }
                Y_AXIS -> {
                    nodePoint = node.point.y
                    pointPlusDistance = query.y + lastDistance
                }
                else -> {
                    nodePoint = node.point.z
                    pointPlusDistance = query.z + lastDistance
                }
            }

            val lineIntersectCube = pointPlusDistance >= nodePoint

            if (lineIntersectCube)
                searchPointsWithinCertainDistance(
                        distance = distance,
                        query = query,
                        node = greater,
                        results = results,
                        examined = examined
                )
        }
    }
}