Skip to content

Instantly share code, notes, and snippets.

@zackthehuman
Created June 9, 2014 13:32
Show Gist options
  • Save zackthehuman/00bb25d3170add090f45 to your computer and use it in GitHub Desktop.
Save zackthehuman/00bb25d3170add090f45 to your computer and use it in GitHub Desktop.
Working generic quadtree written in Scala
case class Identifier(id: Long)
case class Rectangle(x: Int, y: Int, width: Int, height: Int)
object QuadTree {
private val MAX_OBJECTS = 2
private val MAX_LEVELS = 5
private val PARENT_NODE = -1
private val TOP_RIGHT = 0
private val TOP_LEFT = 1
private val BOTTOM_LEFT = 2
private val BOTTOM_RIGHT = 3
}
class QuadTree[T](level: Int, bounds: Rectangle) {
private var objects = List[(T, Rectangle)]()
private var nodes = Array[Option[QuadTree[T]]](None, None, None, None)
def clear() {
objects = List[(T, Rectangle)]()
nodes = nodes.map {
case Some(node) => {
node.clear
None
}
case _ => {
None
}
}
}
private def subdivide() {
val subWidth = bounds.width / 2
val subHeight = bounds.height / 2
val x = bounds.x
val y = bounds.y
nodes(QuadTree.TOP_RIGHT) = Some(new QuadTree(level + 1, new Rectangle(x + subWidth, y, subWidth, subHeight)))
nodes(QuadTree.TOP_LEFT) = Some(new QuadTree(level + 1, new Rectangle(x, y, subWidth, subHeight)))
nodes(QuadTree.BOTTOM_LEFT) = Some(new QuadTree(level + 1, new Rectangle(x, y + subHeight, subWidth, subHeight)))
nodes(QuadTree.BOTTOM_RIGHT) = Some(new QuadTree(level + 1, new Rectangle(x + subWidth, y + subHeight, subWidth, subHeight)))
}
private def getIndex(rect: Rectangle): Int = {
val horizontalMidpoint = bounds.x + bounds.width / 2
val verticalMidpoint = bounds.y + bounds.height / 2
val isTopQuadrant = (rect.y < verticalMidpoint) && ((rect.y + rect.height) < verticalMidpoint)
val isBottomQuadrant = rect.y > verticalMidpoint
var quadrant = QuadTree.PARENT_NODE
if(rect.x < horizontalMidpoint && ((rect.x + rect.width) < horizontalMidpoint)) {
if(isTopQuadrant) {
quadrant = QuadTree.TOP_LEFT
} else if(isBottomQuadrant) {
quadrant = QuadTree.BOTTOM_LEFT
}
} else if(rect.x > horizontalMidpoint) {
if(isTopQuadrant) {
quadrant = QuadTree.TOP_RIGHT
} else if(isBottomQuadrant) {
quadrant = QuadTree.BOTTOM_RIGHT
}
}
quadrant
}
def insert(item: T, itemBounds: Rectangle) {
nodes(0) match {
case Some(node) => {
val index = getIndex(itemBounds)
if(index != QuadTree.PARENT_NODE) {
nodes(index) match {
case Some(node) => node.insert(item, itemBounds)
case _ => throw new RuntimeException("Something went terribly wrong!")
}
return
}
}
case _ => { /* We do nothing on purpose */ }
}
objects = objects :+ (item, itemBounds)
if(objects.size > QuadTree.MAX_OBJECTS && level < QuadTree.MAX_LEVELS) {
nodes(0) match {
case None => {
subdivide()
}
case _ => { /* We do nothing here on purpose */ }
}
var i = 0
while(i < objects.size) {
val itemToReinsert = objects(i)
val index = getIndex(itemToReinsert._2)
if(index != QuadTree.PARENT_NODE) {
nodes(index) match {
case Some(node) => {
objects = objects.filterNot(elm => elm == itemToReinsert)
node.insert(itemToReinsert._1, itemToReinsert._2)
}
case _ => throw new RuntimeException("Something went terribly wrong after subdividing!")
}
} else {
i += 1
}
}
}
}
def query(area: Rectangle): List[(T, Rectangle)] = {
val index = getIndex(area)
var queryResults = List[(T, Rectangle)]()
if(index != QuadTree.PARENT_NODE) {
nodes(index) match {
case Some(node) => {
queryResults = queryResults ++ node.query(area)
}
case _ => { /* Nothing to do in this case */ }
}
}
queryResults = queryResults ++ objects
queryResults
}
override def toString(): String = {
val indent = (for { a <- 0 until level + 1 } yield "\t").mkString("")
val x = bounds.x
val y = bounds.y
val width = bounds.width
val height = bounds.height
return s"Node: ($x, $y, $width, $height) (" + objects.mkString(", ") + ")" +
nodes.flatMap {
case Some(node) => {
Some("\n" + indent + node.toString())
}
case _ => {
None
}
}.mkString("")
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment