Quadtree

Author: timaktimak

QuadTree/QuadTree.playground/Contents.swift

@@ -0,0 +1,14 @@
+
+import Foundation
+
+let tree = QuadTree(rect: Rect(origin: Point(0, 0), size: Size(xLength: 10, yLength: 10)))
+
+for _ in 0..<40 {
+  let randomX = Double(arc4random_uniform(100)) / 10
+  let randomY = Double(arc4random_uniform(100)) / 10
+  let point = Point(randomX, randomY)
+  tree.add(point: point)
+}
+
+print(tree)
+print(tree.points(inRect: Rect(origin: Point(1, 1), size: Size(xLength: 5, yLength: 5))))

QuadTree/QuadTree.playground/Sources/QuadTree.swift

@@ -0,0 +1,290 @@
+
+public struct Point {
+  let x: Double
+  let y: Double
+  
+  public init(_ x: Double, _ y: Double) {
+    self.x = x
+    self.y = y
+  }
+}
+
+extension Point: CustomStringConvertible {
+  public var description: String {
+    return "Point(\(x), \(y))"
+  }
+}
+
+public struct Size: CustomStringConvertible {
+  var xLength: Double
+  var yLength: Double
+  
+  public init(xLength: Double, yLength: Double) {
+    precondition(xLength >= 0, "xLength can not be negative")
+    precondition(yLength >= 0, "yLength can not be negative")
+    self.xLength = xLength
+    self.yLength = yLength
+  }
+  
+  var half: Size {
+    return Size(xLength: xLength / 2, yLength: yLength / 2)
+  }
+  
+  public var description: String {
+    return "Size(\(xLength), \(yLength))"
+  }
+}
+
+public struct Rect {
+  // left top vertice
+  var origin: Point
+  var size: Size
+  
+  public init(origin: Point, size: Size) {
+    self.origin = origin
+    self.size = size
+  }
+  
+  var minX: Double {
+    return origin.x
+  }
+  
+  var minY: Double {
+    return origin.y
+  }
+  
+  var maxX: Double {
+    return origin.x + size.xLength
+  }
+  
+  var maxY: Double {
+    return origin.y + size.yLength
+  }
+  
+  func containts(point: Point) -> Bool {
+    return (minX <= point.x && point.x <= maxX) &&
+      (minY <= point.y && point.y <= maxY)
+  }
+  
+  var leftTopRect: Rect {
+    return Rect(origin: origin, size: size.half)
+  }
+  
+  var leftBottomRect: Rect {
+    return Rect(origin: Point(origin.x, origin.y + size.half.yLength), size: size.half)
+  }
+  
+  var rightTopRect: Rect {
+    return Rect(origin: Point(origin.x + size.half.xLength, origin.y), size: size.half)
+  }
+  
+  var rightBottomRect: Rect {
+    return Rect(origin: Point(origin.x + size.half.xLength, origin.y + size.half.yLength), size: size.half)
+  }
+  
+  func intersects(rect: Rect) -> Bool {
+
+    func lineSegmentsIntersect(lStart: Double, lEnd: Double, rStart: Double, rEnd: Double) -> Bool {
+      return max(lStart, rStart) <= min(lEnd, rEnd)
+    }
+    // to intersect, both horizontal and vertical projections need to intersect
+    // horizontal
+    if !lineSegmentsIntersect(lStart: minX, lEnd: maxX, rStart: rect.minX, rEnd: rect.maxX) {
+      return false
+    }
+    
+    // vertical
+    return lineSegmentsIntersect(lStart: minY, lEnd: maxY, rStart: rect.minY, rEnd: rect.maxY)
+  }
+}
+
+extension Rect: CustomStringConvertible {
+  public var description: String {
+    return "Rect(\(origin), \(size))"
+  }
+}
+
+protocol PointsContainter {
+  func add(point: Point) -> Bool
+  func points(inRect rect: Rect) -> [Point]
+}
+
+class QuadTreeNode {
+  
+  enum NodeType {
+    case leaf
+    case `internal`(children: Children)
+  }
+  
+  struct Children: Sequence {
+    let leftTop: QuadTreeNode
+    let leftBottom: QuadTreeNode
+    let rightTop: QuadTreeNode
+    let rightBottom: QuadTreeNode
+    
+    init(parentNode: QuadTreeNode) {
+      leftTop = QuadTreeNode(rect: parentNode.rect.leftTopRect)
+      leftBottom = QuadTreeNode(rect: parentNode.rect.leftBottomRect)
+      rightTop = QuadTreeNode(rect: parentNode.rect.rightTopRect)
+      rightBottom = QuadTreeNode(rect: parentNode.rect.rightBottomRect)
+    }
+    
+    struct ChildrenIterator: IteratorProtocol {
+      private var index = 0
+      private let children: Children
+      
+      init(children: Children) {
+        self.children = children
+      }
+      
+      mutating func next() -> QuadTreeNode? {
+        
+        defer { index += 1 }
+        
+        switch index {
+        case 0: return children.leftTop
+        case 1: return children.leftBottom
+        case 2: return children.rightTop
+        case 3: return children.rightBottom
+        default: return nil
+        }
+      }
+    }
+    
+    public func makeIterator() -> ChildrenIterator {
+      return ChildrenIterator(children: self)
+    }
+  }
+  
+  var points: [Point] = []
+  let rect: Rect
+  var type: NodeType = .leaf
+  
+  static let maxPointCapacity = 3
+  
+  init(rect: Rect) {
+    self.rect = rect
+  }
+  
+  var recursiveDescription: String {
+    return recursiveDescription(withTabCount: 0)
+  }
+  
+  private func recursiveDescription(withTabCount count: Int) -> String {
+    let indent = String(repeating: "\t", count: count)
+    var result = "\(indent)" + description + "\n"
+    switch type {
+    case .internal(let children):
+      for child in children {
+        result += child.recursiveDescription(withTabCount: count + 1)
+      }
+    default:
+      break
+    }
+    return result
+  }
+}
+
+extension QuadTreeNode: PointsContainter {
+  
+  @discardableResult
+  func add(point: Point) -> Bool {
+    if !rect.containts(point: point) {
+      return false
+    }
+    
+    switch type {
+    case .internal(let children):
+      // pass the point to one of the children
+      for child in children {
+        if child.add(point: point) {
+          return true
+        }
+      }
+      
+      fatalError("rect.containts evaluted to true, but none of the children added the point")
+    case .leaf:
+      points.append(point)
+      // if the max capacity was reached, become an internal node
+      if points.count == QuadTreeNode.maxPointCapacity {
+        subdivide()
+      }
+    }
+    return true
+  }
+  
+  private func subdivide() {
+    switch type {
+    case .leaf:
+      type = .internal(children: Children(parentNode: self))
+    case .internal:
+      preconditionFailure("Calling subdivide on an internal node")
+    }
+  }
+  
+  func points(inRect rect: Rect) -> [Point] {
+    
+    // if the node's rect and the given rect don't intersect, return an empty array,
+    // because there can't be any points that lie the node's (or its children's) rect and
+    // in the given rect
+    if !self.rect.intersects(rect: rect) {
+      return []
+    }
+    
+    var result: [Point] = []
+    
+    // collect the node's points that lie in the rect
+    for point in points {
+      if rect.containts(point: point) {
+        result.append(point)
+      }
+    }
+    
+    switch type {
+    case .leaf:
+      break
+    case .internal(let children):
+      // recursively add children's points that lie in the rect
+      for childNode in children {
+        result.append(contentsOf: childNode.points(inRect: rect))
+      }
+    }
+    
+    return result
+  }
+}
+
+extension QuadTreeNode: CustomStringConvertible {
+  var description: String {
+    switch type {
+    case .leaf:
+      return "leaf \(rect) Points: \(points)"
+    case .internal:
+      return "parent \(rect) Points: \(points)"
+    }
+  }
+}
+
+public class QuadTree: PointsContainter {
+  
+  let root: QuadTreeNode
+  
+  public init(rect: Rect) {
+    self.root = QuadTreeNode(rect: rect)
+  }
+  
+  @discardableResult
+  public func add(point: Point) -> Bool {
+    return root.add(point: point)
+  }
+  
+  public func points(inRect rect: Rect) -> [Point] {
+    return root.points(inRect: rect)
+  }
+}
+
+extension QuadTree: CustomStringConvertible {
+  public var description: String {
+    return "Quad tree\n" + root.recursiveDescription
+  }
+}

QuadTree/QuadTree.playground/contents.xcplayground

@@ -0,0 +1,4 @@
+<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
+<playground version='5.0' target-platform='ios'>
+    <timeline fileName='timeline.xctimeline'/>
+</playground>