Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
import Foundation
import MapKit
extension MKMapPoint {
static var nyc: MKMapPoint {
return MKMapPoint(.nyc)
}
}
extension MKMapRect {
var northWestPoint: MKMapPoint {
MKMapPoint(x: minX, y: maxY)
}
var northEastPoint: MKMapPoint {
MKMapPoint(x: maxX, y: maxY)
}
var southWestPoint: MKMapPoint {
MKMapPoint(x: minX, y: minY)
}
var southEastPoint: MKMapPoint {
MKMapPoint(x: maxX, y: minY)
}
var center: MKMapPoint {
get {
MKMapPoint(x: midX, y: midY)
}
set {
self.origin = MKMapPoint(x: newValue.x - width/2, y: newValue.y - height/2)
}
}
var quadrants: (northWest: MKMapRect, northEast: MKMapRect, southWest: MKMapRect, southEast: MKMapRect) {
var north: MKMapRect = .null
var south: MKMapRect = .null
MKMapRectDivide(self, &north, &south, height/2, .minYEdge)
var northWest: MKMapRect = .null
var northEast: MKMapRect = .null
var southWest: MKMapRect = .null
var southEast: MKMapRect = .null
MKMapRectDivide(north, &northWest, &northEast, width/2, .minXEdge)
MKMapRectDivide(south, &southWest, &southEast, width/2, .minXEdge)
return (northWest, northEast, southWest, southEast)
}
}
protocol MapPointHaving {
var mapPoint: MKMapPoint { get }
}
final class QuadTreeNode<Element: MapPointHaving & Identifiable> {
enum Children {
case undivided([Element])
case divided(northWest: QuadTreeNode, northEast: QuadTreeNode, southWest: QuadTreeNode, southEast: QuadTreeNode)
var areDivided: Bool {
if case .divided = self {
return true
}
return false
}
mutating func divideIfNeeded(rect: MKMapRect) -> (northWest: QuadTreeNode, northEast: QuadTreeNode, southWest: QuadTreeNode, southEast: QuadTreeNode) {
if case .undivided = self {
let quadrants = rect.quadrants
let northWest = QuadTreeNode(rect: quadrants.northWest)
let northEast = QuadTreeNode(rect: quadrants.northEast)
let southWest = QuadTreeNode(rect: quadrants.southWest)
let southEast = QuadTreeNode(rect: quadrants.southEast)
self = .divided(northWest: northWest, northEast: northEast, southWest: southWest, southEast: southEast)
}
return assumingDivided()
}
func ifDivided() -> (northWest: QuadTreeNode, northEast: QuadTreeNode, southWest: QuadTreeNode, southEast: QuadTreeNode)? {
if case let .divided(northWest, northEast, southWest, southEast) = self {
return (northWest, northEast, southWest, southEast)
} else {
return nil
}
}
func assumingDivided() -> (northWest: QuadTreeNode, northEast: QuadTreeNode, southWest: QuadTreeNode, southEast: QuadTreeNode) {
if let divided = ifDivided() {
return divided
} else {
fatalError("Tried to assume that the QuadTree's Children was divided, but it was not")
}
}
}
var children: Children = .undivided([])
var rect: MKMapRect
let bucketSize: Int = 8
init(rect: MKMapRect) {
self.rect = rect
}
@discardableResult
func insert(_ value: Element) -> Bool {
guard self.rect.contains(value.mapPoint) else {
return false
}
switch self.children {
case let .undivided(points):
if points.count < bucketSize {
self.children = .undivided(points + [value])
return true
}
let quadrants = self.children.divideIfNeeded(rect: rect)
for point in points {
quadrants.northWest.insert(point)
quadrants.northEast.insert(point)
quadrants.southWest.insert(point)
quadrants.southEast.insert(point)
}
// could this potentially result in one child having more than bucketSize children?
return quadrants.northWest.insert(value) || quadrants.northEast.insert(value) || quadrants.southWest.insert(value) || quadrants.southEast.insert(value)
case let .divided(northWest, northEast, southWest, southEast):
return northWest.insert(value) || northEast.insert(value) || southWest.insert(value) || southEast.insert(value)
}
}
func elements(in rect: MKMapRect) -> [Element] {
guard self.rect.intersects(rect) else {
return []
}
switch self.children {
case let .undivided(points):
return points.filter({ rect.contains($0.mapPoint) })
case let .divided(northWest, northEast, southWest, southEast):
return northWest.elements(in: rect)
+ northEast.elements(in: rect)
+ southWest.elements(in: rect)
+ southEast.elements(in: rect)
}
}
func closestPlaces(to point: MKMapPoint, maxCount: Int = 20, closerThan distanceInMeters: Double) -> [RelativePlace<Element>] {
if !rect.contains(point)
&& rect.northWestPoint.distance(to: point) > distanceInMeters
&& rect.northEastPoint.distance(to: point) > distanceInMeters
&& rect.southWestPoint.distance(to: point) > distanceInMeters
&& rect.southEastPoint.distance(to: point) > distanceInMeters {
return []
}
let maxDistance = distanceInMeters
// do we need this??
// var maxDistance: Double {
// if points.count >= bucketSize, let farthest = points.max(by: { $0.distance < $1.distance }) {
// return farthest.distance
// } else {
// return .greatestFiniteMagnitude
// }
// }
switch self.children {
case let .undivided(points):
return points
.map({ RelativePlace(place: $0, relativeTo: point) })
.smallest(maxCount)
case let .divided(nw, ne, sw, se):
return (nw.closestPlaces(to: point, maxCount: maxCount, closerThan: maxDistance)
+ ne.closestPlaces(to: point, maxCount: maxCount, closerThan: maxDistance)
+ sw.closestPlaces(to: point, maxCount: maxCount, closerThan: maxDistance)
+ se.closestPlaces(to: point, maxCount: maxCount, closerThan: maxDistance))
.smallest(maxCount)
}
}
}
// extension QuadTreeNode where Element == PlaceAnnotation {
// func update(_ place: Place) {
// //check which quadrant it should go in
// switch self.children {
// case var .undivided(points):
// if let index = points.firstIndex(where: { place.id == $0.id }) {
// let annotation = PlaceAnnotation(place: place)
// points[index] = annotation
// self.children = .undivided(points)
// }
// case let .divided(nw, ne, sw, se):
// ne.update(place)
// nw.update(place)
// se.update(place)
// sw.update(place)
// }
// }
//
// func delete(_ place: Place) {
// switch self.children {
// case var .undivided(points):
// if let index = points.firstIndex(where: { place.id == $0.id }) {
// points.remove(at: index)
// self.children = .undivided(points)
// }
// case let .divided(nw, ne, sw, se):
// ne.delete(place)
// nw.delete(place)
// se.delete(place)
// sw.delete(place)
// }
// }
// }
//
extension QuadTreeNode: Collection {
/// Traversal order - self, northwest, northeast, southwest, southeast
struct Index: Comparable {
enum Branch: Comparable {
case nw, ne, sw, se
var sortOrder: Int {
switch self {
case .nw:
return 0
case .ne:
return 1
case .sw:
return 2
case .se:
return 3
}
}
static func < (lhs: Branch, rhs: Branch) -> Bool {
lhs.sortOrder < rhs.sortOrder
}
}
struct NodeIndex {
// should nodes be weak?
let branchPath: [(parent: QuadTreeNode, direction: Branch)]
let currentNode: QuadTreeNode
func nextNonEmptyUndividedNode() -> NodeIndex? {
return sequence(first: self, next: { $0.nextUndividedNode() })
.dropFirst()
.first(where: {
if case let .undivided(points) = $0.currentNode.children, !points.isEmpty {
return true
} else {
return false
}
})
}
func nextUndividedNode() -> NodeIndex? {
switch branchPath.last {
case let (parent, .nw)?:
var path = Array(branchPath.dropLast()) + [(parent, .ne)]
var node = parent.child(in: .ne)!
while let quadrants = node.children.ifDivided() {
path.append((node, .nw))
node = quadrants.northWest
}
return NodeIndex(branchPath: path, currentNode: node)
case let (parent, .ne)?:
var path = Array(branchPath.dropLast()) + [(parent, .sw)]
var node = parent.child(in: .sw)!
while let quadrants = node.children.ifDivided() {
path.append((node, .nw))
node = quadrants.northWest
}
return NodeIndex(branchPath: path, currentNode: node)
case let (parent, .sw)?:
var path = Array(branchPath.dropLast()) + [(parent, .se)]
var node = parent.child(in: .se)!
while let quadrants = node.children.ifDivided() {
path.append((node, .nw))
node = quadrants.northWest
}
return NodeIndex(branchPath: path, currentNode: node)
case let (parent, .se)?:
return NodeIndex(branchPath: Array(branchPath.dropLast()), currentNode: parent)
.nextUndividedNode()
case nil:
return nil
}
}
}
var nodeIndex: NodeIndex
let arrayIndex: Int
init(branchPath: [(parent: QuadTreeNode, direction: Branch)], currentNode: QuadTreeNode, arrayIndex: Int) {
self.nodeIndex = NodeIndex(branchPath: branchPath, currentNode: currentNode)
self.arrayIndex = arrayIndex
}
var branchPath: [(parent: QuadTreeNode, direction: Branch)] {
nodeIndex.branchPath
}
var currentNode: QuadTreeNode {
nodeIndex.currentNode
}
static func < (lhs: Index, rhs: Index) -> Bool {
let lhsDirections = lhs.branchPath.lazy.map({ $0.direction })
let rhsDirections = rhs.branchPath.lazy.map({ $0.direction })
if lhsDirections.elementsEqual(rhsDirections) {
return lhs.arrayIndex < rhs.arrayIndex
} else {
return lhsDirections.lexicographicallyPrecedes(rhsDirections)
}
}
static func == (lhs: Index, rhs: Index) -> Bool {
let lhsDirections = lhs.branchPath.lazy.map({ $0.direction })
let rhsDirections = rhs.branchPath.lazy.map({ $0.direction })
return lhsDirections.elementsEqual(rhsDirections) && lhs.arrayIndex == rhs.arrayIndex
}
}
var startIndex: Index {
var current = self
var path: [(QuadTreeNode, Index.Branch)] = []
while case let .divided(nw, _, _, _) = current.children {
path.append((current, .nw))
current = nw
}
guard let index = Index.NodeIndex(branchPath: path, currentNode: current).nextNonEmptyUndividedNode() else { return endIndex }
return Index(branchPath: index.branchPath, currentNode: index.currentNode, arrayIndex: 0)
}
var endIndex: Index {
var current = self
var path: [(QuadTreeNode, Index.Branch)] = []
while case let .divided(_, _, _, se) = current.children {
path.append((current, .se))
current = se
}
return Index(branchPath: path, currentNode: current, arrayIndex: Int.max)
}
func index(after i: Index) -> Index {
guard case let .undivided(points) = i.currentNode.children else {
fatalError("Every index should point to a leaf node")
}
if i.arrayIndex < points.count - 1 {
return Index(branchPath: i.branchPath, currentNode: i.currentNode, arrayIndex: i.arrayIndex + 1)
}
guard let j = i.nodeIndex.nextNonEmptyUndividedNode() else {
return endIndex
}
return Index(branchPath: j.branchPath, currentNode: j.currentNode, arrayIndex: 0)
}
subscript(position: Index) -> Element {
if case let .undivided(points) = position.currentNode.children {
return points[position.arrayIndex]
} else {
fatalError("out of bounds!")
}
}
func node<C: Collection>(for path: C) -> QuadTreeNode? where C.Element == Index.Branch {
guard let firstBranch = path.first else { return self }
//invalid path, maybe should crash by changing ? to !
return self.child(in: firstBranch)?.node(for: path.dropFirst())
}
func child(in direction: Index.Branch) -> QuadTreeNode? {
if let quadrants = self.children.ifDivided() {
switch direction {
case .nw:
return quadrants.northWest
case .ne:
return quadrants.northEast
case .sw:
return quadrants.southWest
case .se:
return quadrants.southEast
}
}
return nil
}
}
extension QuadTreeNode {
func asArray() -> [Element] {
switch self.children {
case let .undivided(points):
return points
case let .divided(nw, ne, sw, se):
return nw.asArray() + ne.asArray() + sw.asArray() + se.asArray()
}
}
}
extension Sequence {
func smallest(_ n: Int, usingTest areInIncreasingOrder: (Element, Element) -> Bool) -> [Element] {
var result = self.prefix(n).sorted(by: areInIncreasingOrder)
for e in self.dropFirst(n) {
if let last = result.last, areInIncreasingOrder(last, e) { continue }
if let insertionIndex = result.firstIndex(where: { areInIncreasingOrder(e, $0) }) {
result.insert(e, at: insertionIndex)
result.removeLast()
}
}
return result
}
}
extension Sequence where Element: Comparable {
func smallest(_ n: Int) -> [Element] {
self.smallest(n, usingTest: <)
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment