Implemented drop/take/slice/splitAt/dropRight/takeRight for

TreeMap/TreeSet by splitting the underlying RedBlack tree. This
makes the operation O(log n) instead of O(n) and allows more
structural sharing.

1 files changed, 23 insertions, 0 deletions

diff --git a/src/library/scala/collection/immutable/TreeSet.scala b/src/library/scala/collection/immutable/TreeSet.scala
index b969ecc0e8..dfaffcd581 100644
--- a/src/library/scala/collection/immutable/TreeSet.scala
+++ b/src/library/scala/collection/immutable/TreeSet.scala
@@ -61,6 +61,29 @@ class TreeSet[A](override val size: Int, t: RedBlack[A]#Tree[Unit])
   override def tail = new TreeSet(size - 1, tree.delete(firstKey))
   override def init = new TreeSet(size - 1, tree.delete(lastKey))
 
+  override def drop(n: Int) = {
+    if (n <= 0) this
+    else if (n >= size) empty
+    else from(tree.nth(n).key)
+  }
+
+  override def take(n: Int) = {
+    if (n <= 0) empty
+    else if (n >= size) this
+    else until(tree.nth(n).key)
+  }
+
+  override def slice(from: Int, until: Int) = {
+    if (until <= from) empty
+    else if (from <= 0) take(until)
+    else if (until >= size) drop(from)
+    else range(tree.nth(from).key, tree.nth(until).key)
+  }
+
+  override def dropRight(n: Int) = take(size - n)
+  override def takeRight(n: Int) = drop(size - n)
+  override def splitAt(n: Int) = (take(n), drop(n))
+
   def isSmaller(x: A, y: A) = compare(x,y) < 0
 
   def this()(implicit ordering: Ordering[A]) = this(0, null)(ordering)