78 lines
1.7 KiB
Scala
78 lines
1.7 KiB
Scala
package org.mimirdb.pip.lib
|
|
|
|
class KDTree[I: Ordered, V](dimensions: Int)
|
|
{
|
|
type Key = Array[I]
|
|
var root: Option[Node] = None
|
|
|
|
def insert(position: Key, value: V): Unit =
|
|
{
|
|
assert(position.size == dimensions)
|
|
root = root match {
|
|
case Some(node) => node.insert(position, value)
|
|
case None => Leaf(position, value, 0)
|
|
}
|
|
}
|
|
|
|
|
|
sealed trait Node
|
|
{
|
|
def insert(position: Key, value: V): Node
|
|
}
|
|
|
|
class Inner(
|
|
split: I,
|
|
idx: Int,
|
|
var left: Option[Node] = None,
|
|
var right: Option[Node] = None
|
|
)
|
|
{
|
|
|
|
def insert(position: Key, value: V): Node =
|
|
{
|
|
if(position(idx).compare(split) >= 0){
|
|
// position(idx) >= split
|
|
if(right.isEmpty){
|
|
right = Some(Leaf(position, value, idx+1 % dimensions))
|
|
} else {
|
|
right = Some(right.get.insert(position, value))
|
|
}
|
|
} else {
|
|
// position(idx) < split
|
|
if(left.isEmpty){
|
|
left = Some(Leaf(position, value, idx+1 % dimensions))
|
|
} else {
|
|
left = Some(left.get.insert(position, value))
|
|
}
|
|
|
|
}
|
|
return this
|
|
}
|
|
}
|
|
|
|
|
|
class Leaf(position: Key, value: V, idx: Int)
|
|
{
|
|
def insert(otherPosition: Key, otherValue: V): Node =
|
|
{
|
|
if(position(idx).compare(otherPosition(idx)) >= 0){
|
|
Inner(
|
|
split = position(idx),
|
|
idx = idx,
|
|
left = Some(Leaf(otherPosition, otherValue, idx+1 % dimensions)),
|
|
right = Some(copy(idx = idx+1 % dimensions)),
|
|
)
|
|
} else {
|
|
Inner(
|
|
split = otherPosition(idx),
|
|
idx = idx,
|
|
left = Some(copy(idx = idx+1 % dimensions)),
|
|
right = Some(Leaf(otherPosition, otherValue, idx+1 % dimensions)),
|
|
)
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
}
|