Skip to content

Commit 51ddeb3

Browse files
author
aleksandar
committed
Add mutable tree sets to the standard library.
This implementation is based on AVL trees. The current implementation is contributed by Lucien Pereira. Fixes #4147.
1 parent 178d49d commit 51ddeb3

File tree

11 files changed

+667
-1
lines changed

11 files changed

+667
-1
lines changed

‎.gitignore

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,3 @@
11
*.jar
2+
build
3+
*.*~

‎src/library/scala/collection/generic/MutableSortedSetFactory.scala

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
/* __ *\
2+
** ________ ___ / / ___ Scala API **
3+
** / __/ __// _ | / / / _ | (c) 2003-2011, LAMP/EPFL **
4+
** __\ \/ /__/ __ |/ /__/ __ | http://scala-lang.org/ **
5+
** /____/\___/_/ |_/____/_/ | | **
6+
** |/ **
7+
\* */
8+
9+
package scala.collection
10+
package generic
11+
12+
import scala.collection.mutable.{ Builder, GrowingBuilder }
13+
14+
/**
15+
* @define Coll mutable.SortedSet
16+
* @define coll mutable sorted
17+
*
18+
* @author Lucien Pereira
19+
*
20+
*/
21+
abstract class MutableSortedSetFactory[CC[A] <: mutable.SortedSet[A] with SortedSetLike[A, CC[A]] with mutable.Set[A] with mutable.SetLike[A, CC[A]]] extends SortedSetFactory[CC] {
22+
23+
/**
24+
* mutable.SetBuilder uses '+' which is not a primitive for anything extending mutable.SetLike,
25+
* this causes serious perfomances issues since each time 'elems = elems + x'
26+
* is evaluated elems is cloned (which is O(n)).
27+
*
28+
* Fortunately GrowingBuilder comes to rescue.
29+
*
30+
*/
31+
override def newBuilder[A](implicit ord: Ordering[A]): Builder[A, CC[A]] = new GrowingBuilder[A, CC[A]](empty)
32+
33+
}

‎src/library/scala/collection/mutable/AVLTree.scala

Lines changed: 206 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,206 @@
1+
/* __ *\
2+
** ________ ___ / / ___ Scala API **
3+
** / __/ __// _ | / / / _ | (c) 2003-2011, LAMP/EPFL **
4+
** __\ \/ /__/ __ |/ /__/ __ | http://scala-lang.org/ **
5+
** /____/\___/_/ |_/____/_/ | | **
6+
** |/ **
7+
\* */
8+
9+
package scala.collection
10+
package mutable
11+
12+
import annotation.tailrec
13+
14+
/**
15+
* An immutable AVL Tree implementation used by mutable.TreeSet
16+
*
17+
* @author Lucien Pereira
18+
*
19+
*/
20+
private[mutable] sealed trait AVLTree[+A] extends Serializable {
21+
def balance: Int
22+
23+
def depth: Int
24+
25+
}
26+
27+
private case class Node[A](val data: A, val left: AVLTree[A], val right: AVLTree[A]) extends AVLTree[A] {
28+
override val balance: Int = right.depth - left.depth
29+
30+
override val depth: Int = math.max(left.depth, right.depth) + 1
31+
32+
}
33+
34+
private case object Leaf extends AVLTree[Nothing] {
35+
override val balance: Int = 0
36+
37+
override val depth: Int = -1
38+
39+
}
40+
41+
private[mutable] object AVLTree {
42+
43+
/**
44+
* Returns a new tree containing the given element.
45+
* Thows an IllegalArgumentException if element is already present.
46+
*
47+
*/
48+
def insert[A](value: A, tree: AVLTree[A], ordering: Ordering[A]): AVLTree[A] = {
49+
@tailrec
50+
def insertTC(value: A, tree: AVLTree[A], reassemble: AVLTree[A] => AVLTree[A]): AVLTree[A] = tree match {
51+
case Leaf => reassemble(Node(value, Leaf, Leaf))
52+
53+
case Node(a, left, right) => if (0 == ordering.compare(value, a)) {
54+
throw new IllegalArgumentException()
55+
} else if (-1 == ordering.compare(value, a)) {
56+
insertTC(value, left, x => reassemble(rebalance(Node(a, x, right))))
57+
} else {
58+
insertTC(value, right, x => reassemble(rebalance(Node(a, left, x))))
59+
}
60+
}
61+
62+
insertTC(value, tree, x => rebalance(x))
63+
}
64+
65+
def contains[A](value: A, tree: AVLTree[A], ordering: Ordering[A]): Boolean = tree match {
66+
case Leaf => false
67+
68+
case Node(a, left, right) => if (0 == ordering.compare(value, a)) {
69+
true
70+
} else if (-1 == ordering.compare(value, a)) {
71+
contains(value, left, ordering)
72+
} else {
73+
contains(value, right, ordering)
74+
}
75+
}
76+
77+
/**
78+
* Return a new tree which not contains given element.
79+
*
80+
*/
81+
def remove[A](value: A, tree: AVLTree[A], ordering: Ordering[A]): AVLTree[A] = tree match {
82+
case Leaf => throw new NoSuchElementException()
83+
84+
case Node(a, Leaf, Leaf) => if (0 == ordering.compare(value, a)) {
85+
Leaf
86+
} else {
87+
throw new NoSuchElementException()
88+
}
89+
90+
case Node(a, left, right@Node(_, _, _)) => if (0 == ordering.compare(value, a)) {
91+
val (min, newRight) = removeMin(right)
92+
rebalance(Node(min, left, newRight))
93+
} else if (-1 == ordering.compare(value, a)) {
94+
rebalance(Node(a, remove(value, left, ordering), right))
95+
} else {
96+
rebalance(Node(a, left, remove(value, right, ordering)))
97+
}
98+
99+
case Node(a, left@Node(_, _, _), right) => if (0 == ordering.compare(value, a)) {
100+
val (max, newLeft) = removeMax(left)
101+
rebalance(Node(max, newLeft, right))
102+
} else if (-1 == ordering.compare(value, a)) {
103+
rebalance(Node(a, remove(value, left, ordering), right))
104+
} else {
105+
rebalance(Node(a, left, remove(value, right, ordering)))
106+
}
107+
}
108+
109+
/**
110+
* Return a tuple containing the biggest element of the provided tree
111+
* and a new tree from which this element has been extracted.
112+
*
113+
*/
114+
def removeMax[A](tree: Node[A]): (A, AVLTree[A]) = {
115+
@tailrec
116+
def removeMaxTC(tree: AVLTree[A], assemble: (A, AVLTree[A]) => (A, AVLTree[A])): (A, AVLTree[A]) = tree match {
117+
case Node(a, Leaf, Leaf) => assemble(a, Leaf)
118+
case Node(a, left, Leaf) => assemble(a, left)
119+
case Node(a, left, right) => removeMaxTC(right,
120+
(max: A, avl: AVLTree[A]) => assemble(max, rebalance(Node(a, left, avl))))
121+
case Leaf => sys.error("Should not happen.")
122+
}
123+
124+
removeMaxTC(tree, (a, b) => (a, b))
125+
}
126+
127+
/**
128+
* Return a tuple containing the smallest element of the provided tree
129+
* and a new tree from which this element has been extracted.
130+
*
131+
*/
132+
def removeMin[A](tree: Node[A]): (A, AVLTree[A]) = {
133+
@tailrec
134+
def removeMinTC(tree: AVLTree[A], assemble: (A, AVLTree[A]) => (A, AVLTree[A])): (A, AVLTree[A]) = tree match {
135+
case Node(a, Leaf, Leaf) => assemble(a, Leaf)
136+
case Node(a, Leaf, right) => assemble(a, right)
137+
case Node(a, left, right) => removeMinTC(left,
138+
(min: A, avl: AVLTree[A]) => assemble(min, rebalance(Node(a, avl, right))))
139+
case Leaf => sys.error("Should not happen.")
140+
}
141+
142+
removeMinTC(tree, (a, b) => (a, b))
143+
}
144+
145+
/**
146+
* Returns a bounded stream of elements in the tree.
147+
*
148+
*/
149+
def toStream[A](tree: AVLTree[A], isLeftAcceptable: A => Boolean, isRightAcceptable: A => Boolean): Stream[A] = tree match {
150+
case Leaf => Stream.empty
151+
152+
case Node(a, left, right) => if (isLeftAcceptable(a)) {
153+
if (isRightAcceptable(a)) {
154+
toStream(left, isLeftAcceptable, isRightAcceptable) ++ Stream(a) ++ toStream(right, isLeftAcceptable, isRightAcceptable)
155+
} else {
156+
toStream(left, isLeftAcceptable, isRightAcceptable)
157+
}
158+
} else if (isRightAcceptable(a)) {
159+
toStream(right, isLeftAcceptable, isRightAcceptable)
160+
} else {
161+
Stream.empty
162+
}
163+
}
164+
165+
/**
166+
* Returns a bounded iterator of elements in the tree.
167+
*
168+
*/
169+
def iterator[A](tree: AVLTree[A], isLeftAcceptable: A => Boolean, isRightAcceptable: A => Boolean): Iterator[A] =
170+
toStream(tree, isLeftAcceptable, isRightAcceptable).iterator
171+
172+
def rebalance[A](tree: AVLTree[A]): AVLTree[A] = (tree, tree.balance) match {
173+
case (node@Node(_, left, _), -2) => left.balance match {
174+
case 1 => doubleRightRotation(node)
175+
case _ => rightRotation(node)
176+
}
177+
178+
case (node@Node(_, _, right), 2) => right.balance match {
179+
case -1 => doubleLeftRotation(node)
180+
case _ => leftRotation(node)
181+
}
182+
183+
case _ => tree
184+
}
185+
186+
def leftRotation[A](tree: Node[A]): AVLTree[A] = tree.right match {
187+
case Node(b, left, right) => Node(b, Node(tree.data, tree.left, left), right)
188+
case _ => sys.error("Should not happen.")
189+
}
190+
191+
def rightRotation[A](tree: Node[A]): AVLTree[A] = tree.left match {
192+
case Node(b, left, right) => Node(b, left, Node(tree.data, right, tree.right))
193+
case _ => sys.error("Should not happen.")
194+
}
195+
196+
def doubleLeftRotation[A](tree: Node[A]): AVLTree[A] = tree.right match {
197+
case right@Node(b, l, r) => leftRotation(Node(tree.data, tree.left, rightRotation(right)))
198+
case _ => sys.error("Should not happen.")
199+
}
200+
201+
def doubleRightRotation[A](tree: Node[A]): AVLTree[A] = tree.left match {
202+
case left@Node(b, l, r) => rightRotation(Node(tree.data, leftRotation(left), tree.right))
203+
case _ => sys.error("Should not happen.")
204+
}
205+
206+
}

‎src/library/scala/collection/mutable/SortedSet.scala

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
/* __ *\
2+
** ________ ___ / / ___ Scala API **
3+
** / __/ __// _ | / / / _ | (c) 2003-2011, LAMP/EPFL **
4+
** __\ \/ /__/ __ |/ /__/ __ | http://scala-lang.org/ **
5+
** /____/\___/_/ |_/____/_/ | | **
6+
** |/ **
7+
\* */
8+
9+
package scala.collection
10+
package mutable
11+
12+
import generic._
13+
14+
/**
15+
* Base trait for mutable sorted set.
16+
*
17+
* @define Coll mutable.SortedSet
18+
* @define coll mutable sorted set
19+
*
20+
* @author Lucien Pereira
21+
*
22+
*/
23+
trait SortedSet[A] extends collection.SortedSet[A] with collection.SortedSetLike[A,SortedSet[A]]
24+
with mutable.Set[A] with mutable.SetLike[A, SortedSet[A]] {
25+
26+
/** Needs to be overridden in subclasses. */
27+
override def empty: SortedSet[A] = SortedSet.empty[A]
28+
29+
}
30+
31+
/**
32+
* A template for mutable sorted set companion objects.
33+
*
34+
* @define Coll mutable.SortedSet
35+
* @define coll mutable sorted set
36+
* @define factoryInfo
37+
* This object provides a set of operations needed to create sorted sets of type mutable.SortedSet.
38+
* @define sortedSetCanBuildFromInfo
39+
* Standard `CanBuildFrom` instance for sorted sets.
40+
*
41+
* @author Lucien Pereira
42+
*
43+
*/
44+
object SortedSet extends MutableSortedSetFactory[SortedSet] {
45+
implicit def canBuildFrom[A](implicit ord: Ordering[A]): CanBuildFrom[Coll, A, SortedSet[A]] = new SortedSetCanBuildFrom[A]
46+
47+
def empty[A](implicit ord: Ordering[A]): SortedSet[A] = TreeSet.empty[A]
48+
49+
}

0 commit comments

Comments
 (0)