diff options
Diffstat (limited to 'src')
-rw-r--r-- | src/library/scala/collection/immutable/RedBlackTree.scala | 25 |
1 files changed, 17 insertions, 8 deletions
diff --git a/src/library/scala/collection/immutable/RedBlackTree.scala b/src/library/scala/collection/immutable/RedBlackTree.scala index 5bdba26d02..bb489dd80a 100644 --- a/src/library/scala/collection/immutable/RedBlackTree.scala +++ b/src/library/scala/collection/immutable/RedBlackTree.scala @@ -132,6 +132,15 @@ object RedBlackTree { else if (overwrite || k != tree.key) mkTree(isBlackTree(tree), k, v, tree.left, tree.right) else tree } + private[this] def updNth[A, B, B1 >: B](tree: Tree[A, B], idx: Int, k: A, v: B1, overwrite: Boolean): Tree[A, B1] = if (tree eq null) { + RedTree(k, v, null, null) + } else { + val rank = count(tree.left) + 1 + if (idx < rank) balanceLeft(isBlackTree(tree), tree.key, tree.value, updNth(tree.left, idx, k, v, overwrite), tree.right) + else if (idx > rank) balanceRight(isBlackTree(tree), tree.key, tree.value, tree.left, updNth(tree.right, idx - rank, k, v, overwrite)) + else if (overwrite) mkTree(isBlackTree(tree), k, v, tree.left, tree.right) + else tree + } /* Based on Stefan Kahrs' Haskell version of Okasaki's Red&Black Trees * http://www.cse.unsw.edu.au/~dons/data/RedBlackTree.html */ @@ -249,27 +258,27 @@ object RedBlackTree { else rebalance(tree, newLeft, newRight) } - private[this] def doDrop[A: Ordering, B](tree: Tree[A, B], n: Int): Tree[A, B] = { + private[this] def doDrop[A, B](tree: Tree[A, B], n: Int): Tree[A, B] = { if (n <= 0) return tree if (n >= this.count(tree)) return null val count = this.count(tree.left) if (n > count) return doDrop(tree.right, n - count - 1) val newLeft = doDrop(tree.left, n) if (newLeft eq tree.left) tree - else if (newLeft eq null) upd(tree.right, tree.key, tree.value, false) + else if (newLeft eq null) updNth(tree.right, n - count - 1, tree.key, tree.value, false) else rebalance(tree, newLeft, tree.right) } - private[this] def doTake[A: Ordering, B](tree: Tree[A, B], n: Int): Tree[A, B] = { + private[this] def doTake[A, B](tree: Tree[A, B], n: Int): Tree[A, B] = { if (n <= 0) return null if (n >= this.count(tree)) return tree val count = this.count(tree.left) if (n <= count) return doTake(tree.left, n) val newRight = doTake(tree.right, n - count - 1) if (newRight eq tree.right) tree - else if (newRight eq null) upd(tree.left, tree.key, tree.value, false) + else if (newRight eq null) updNth(tree.left, n, tree.key, tree.value, false) else rebalance(tree, tree.left, newRight) } - private[this] def doSlice[A: Ordering, B](tree: Tree[A, B], from: Int, until: Int): Tree[A, B] = { + private[this] def doSlice[A, B](tree: Tree[A, B], from: Int, until: Int): Tree[A, B] = { if (tree eq null) return null val count = this.count(tree.left) if (from > count) return doSlice(tree.right, from - count - 1, until - count - 1) @@ -277,8 +286,8 @@ object RedBlackTree { val newLeft = doDrop(tree.left, from) val newRight = doTake(tree.right, until - count - 1) if ((newLeft eq tree.left) && (newRight eq tree.right)) tree - else if (newLeft eq null) upd(newRight, tree.key, tree.value, false) - else if (newRight eq null) upd(newLeft, tree.key, tree.value, false) + else if (newLeft eq null) updNth(newRight, from - count - 1, tree.key, tree.value, false) + else if (newRight eq null) updNth(newLeft, until, tree.key, tree.value, false) else rebalance(tree, newLeft, newRight) } @@ -380,7 +389,7 @@ object RedBlackTree { @(inline @getter) final val left: Tree[A, B], @(inline @getter) final val right: Tree[A, B]) extends Serializable { - final val count: Int = 1 + RedBlackTree.count(left) + RedBlackTree.count(right) + @(inline @getter) final val count: Int = 1 + RedBlackTree.count(left) + RedBlackTree.count(right) def black: Tree[A, B] def red: Tree[A, B] } |