-
-
Save yoeluk/70bd27398a31f1a15e16ad4315df4e8d to your computer and use it in GitHub Desktop.
Simple immutable B-tree with Scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
case class BTree[K, V](root: BTree.Node[K, V], order: Int)(implicit keyOrdering: Ordering[K]) { | |
import keyOrdering.mkOrderingOps | |
private type N = BTree.Node[K, V] | |
private type E = BTree.Entry[K, V] | |
def get(key: K): Option[V] = { | |
def aux(node: N): Option[V] = { | |
val mayBeEntry = node.entries.find(_.key == key) | |
if (node.leaf) mayBeEntry.map(_.value) else { | |
mayBeEntry match { | |
case None => aux(node.children(childIndex(node, key))) | |
case Some(entry) => Some(entry.value) | |
} | |
} | |
} | |
aux(root) | |
} | |
def insert(newKey: K, newValue: V): BTree[K, V] = { | |
val t = this.order | |
def splitNode(node: N): (E, N, N) = ( | |
node.entries(t - 1), | |
BTree.Node( | |
entries = node.entries.take(t - 1), | |
children = node.children.take(t) | |
), | |
BTree.Node( | |
entries = node.entries.takeRight(t - 1), | |
children = node.children.takeRight(t) | |
) | |
) | |
def tryInsert(node: N): Either[(E, N, N), N] = { | |
val maybeEntryIndex = node.entries.indexWhere(_.key == newKey) | |
if (maybeEntryIndex > -1) { | |
// Key already exists | |
val newEntry = BTree.Entry(newKey, newValue) | |
val newData = node.entries.updated(maybeEntryIndex, newEntry) | |
Right(node.copy(entries = newData)) | |
} else { | |
val modified = | |
if (node.leaf) { | |
val newEntry = BTree.Entry(newKey, newValue) | |
val newData = (node.entries :+ newEntry).sortBy(_.key) | |
node.copy(entries = newData) | |
} else { | |
val index = childIndex(node, newKey) | |
val child = node.children(index) | |
tryInsert(child) match { | |
case Left((restEntry, splitLhs, splitRhs)) => | |
val (lhs, rhs) = node.children.splitAt(index) | |
node.copy( | |
entries = (node.entries :+ restEntry).sortBy(_.key), | |
children = (lhs :+ splitLhs :+ splitRhs) ++ rhs.tail | |
) | |
case Right(modifiedChild) => | |
node.copy(children = node.children.updated(index, modifiedChild)) | |
} | |
} | |
if (modified.entries.length == 2 * t - 1) Left(splitNode(modified)) | |
else Right(modified) | |
} | |
} | |
this.copy( | |
root = tryInsert(root) match { | |
case Right(node) => node | |
case Left((restEntry, lhs, rhs)) => | |
BTree.Node(Vector(restEntry), Vector(lhs, rhs)) | |
} | |
) | |
} | |
/** For debug purposes */ | |
def renderStructure: String = { | |
def aux(node: N, i: String): String = { | |
val keysString = node.entries.map(x => x.key).mkString(",") | |
if (node.children.nonEmpty) { | |
val childrenString = node.children.map(x => aux(x, i + " ")).mkString("\n") | |
s"$i$keysString\n$childrenString" | |
} else { | |
s"$i$keysString" | |
} | |
} | |
aux(root, "") | |
} | |
private def childIndex(node: N, key: K): Int = | |
node.entries.lastIndexWhere(key > _.key) + 1 | |
} | |
object BTree { | |
case class Entry[K, V](key: K, value: V) | |
case class Node[K: Ordering, V](entries: Vector[Entry[K, V]], children: Vector[Node[K, V]]) { | |
def leaf: Boolean = children.isEmpty | |
} | |
def empty[K: Ordering, V](order: Int): BTree[K, V] = | |
BTree(Node[K, V](Vector.empty, Vector.empty), order) | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment