Skip to content

Instantly share code, notes, and snippets.

@yellowflash
Last active June 27, 2018 07:56
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
Star You must be signed in to star a gist
Save yellowflash/3e7a1b04b5889999b2af5ff7f736e654 to your computer and use it in GitHub Desktop.
Spatial Index with geo hash
import scala.annotation.tailrec
import scala.collection.mutable
case class Point(x: Int, y: Int) {
def geoHash(precision: Int): BigInt = {
@tailrec
def hash(minX: Int, maxX: Int, minY: Int, maxY: Int, prec: Int, sofar: BigInt) : BigInt = {
if(prec == 0) return sofar;
val quadrant = if((minX + maxX)/2 < x) {
if((minY + maxY)/2 < y) 0 else 1
} else {
if((minY + maxY)/2 < y) 2 else 3
}
quadrant match {
case 0 => hash(minX, (minX + maxX)/2, minY, (minY + maxY)/2, prec - 1, (sofar << 2) + quadrant)
case 1 => hash(minX, (minX + maxX)/2, (minY + maxY)/2, maxY, prec - 1, (sofar << 2) + quadrant)
case 2 => hash((minX + maxX)/2, maxX, minY, (minY + maxY)/2, prec - 1, (sofar << 2) + quadrant)
case 3 => hash((minX + maxX)/2, maxX, (minY + maxY)/2, maxY, prec - 1, (sofar << 2) + quadrant)
}
}
hash(Int.MinValue, Int.MaxValue, Int.MinValue, Int.MaxValue, precision, 0)
}
}
class SpatialMap[V](precision: Int) {
val values = mutable.SortedMap[BigInt, List[V]]()
def add(point: Point, value: V) = {
val geohash = point.geoHash(precision)
values += (geohash -> (value::values.getOrElse(geohash, List.empty)))
}
def nearby(point: Point, qprec: Int): Iterator[V] = {
assert(qprec < precision)
val start = point.geoHash(qprec) << ((precision - qprec) * 2)
val end = shiftPadOne(point.geoHash(qprec), (precision - qprec) * 2)
//TODO: This might not give you the square in which the point is at center.
//We need to do 4 such lookups, which is just trivial.
values.iteratorFrom(start).takeWhile(_._1 <= end).flatMap(_._2)
}
@tailrec
private def shiftPadOne(value: BigInt, bits: Int): BigInt = if(bits == 0) value else shiftPadOne((value << 1) | 1, bits - 1)
}
object SpatialMap {
def main(args: Array[String]): Unit = {
val index = new SpatialMap[String](32)
index.add(Point(1,1), "a")
index.add(Point(200, 200), "b")
index.add(Point(3000, 3000), "c")
index.add(Point(40000, 40000), "d")
index.add(Point(500000, 500000), "e")
index.add(Point(6000000, 6000000), "f")
index.add(Point(6000010, 6000010), "f'")
index.add(Point(70000000, 70000000), "g")
println(index.nearby(Point(1,1), 25).toList)
println(index.nearby(Point(1,1), 21).toList)
println(index.nearby(Point(1,1), 20).toList)
println(index.nearby(Point(1,1), 15).toList)
println(index.nearby(Point(6000000,6000000), 21).toList)
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment