Bersaelor / KDTree

Swift implementation of a k-dimensional binary space partitioning tree.
MIT License
153 stars 26 forks source link

Have a conditional in the kNearest method #30

Closed chriseidhof closed 6 years ago

chriseidhof commented 6 years ago

Is it possible to have a conditional in the kNearest method?

I'm thinking of something like this:

extension KDTree {

    /// Returns the k nearest `KDTreePoint`s to the search point,
    ///
    /// - Complexity: O(log N).
    public func nearestK(_ number: Int, to searchElement: Element) -> [Element] {
        var neighbours: Neighbours = Neighbours(goalNumber: number)
        self.nearestK(to: searchElement, where: { _ in true }, bestValues: &neighbours)
        return neighbours.nearestValues.map { $0.point as! Element }
    }

    public func nearestK(_ number: Int, to searchElement: Element, where condition: (Element) -> Bool) -> [Element] {
        var neighbours: Neighbours = Neighbours(goalNumber: number)
        self.nearestK(to: searchElement, where: condition, bestValues: &neighbours)
        return neighbours.nearestValues.map { $0.point as! Element }
    }

    fileprivate func nearestK(to searchElement: Element, where condition: (Element) -> Bool, bestValues: inout Neighbours) {
        switch self {
        case let .node(left, value, dim, right):
            let dimensionDifference = value.kdDimension(dim) - searchElement.kdDimension(dim)
            let isLeftOfValue = dimensionDifference > 0

            //check the best estimate side
            let closerSubtree = isLeftOfValue ? left : right
            closerSubtree.nearestK(to: searchElement, where: condition, bestValues: &bestValues)

            if condition(value) {
                //check the nodes value
                let currentDistance = value.squaredDistance(to: searchElement)
                bestValues.append(value, distance: currentDistance)
            }

            //if the bestDistance so far intersects the hyperplane at the other side of this value
            //there could be points in the other subtree
            if dimensionDifference*dimensionDifference < bestValues.biggestDistance || !bestValues.full {
                let otherSubtree = isLeftOfValue ? right : left
                otherSubtree.nearestK(to: searchElement, where: condition, bestValues: &bestValues)
            }
        case .leaf: break
        }
    }
}

It compiles and works, but I'm not sure if I'm breaking any invariants here...

Bersaelor commented 6 years ago

I think this should be fine. If a lot of points are rejected, the algorithm will eventually explore all subtree's, as without setting 'bestNewDistance' it will always check on the other side.

You might loose some performance because more subtree's will be checked, but it will still find the point closest to your start that fits the condition. Whether the points that the space is split at are part of the actual nearest-check or not doesn't change the outcome.

Bersaelor commented 6 years ago

You can make a PR where the condition is nil or { _ in true} by default if you like 🙃 (If you do, please do for both nearest and the nearestK algorithm implementations, I split them since the nearest is easier to reason for without the neighbours-collection need for k neighbours)

chriseidhof commented 6 years ago

See #31