#### /core/src/main/scala/scalaz/Heap.scala

http://github.com/scalaz/scalaz
Scala | 436 lines | 299 code | 88 blank | 49 comment | 47 complexity | f5d10cf9c7c5936a277337111338c490 MD5 | raw file
```  1package scalaz
2
3import std.tuple._
4import EphemeralStream.{EStream, emptyEphemeralStream, ##::}
5
6
7/**An efficient, asymptotically optimal, implementation of priority queues
8 * extended with support for efficient size.
9 *
10 * The implementation of 'Heap' is based on bootstrapped skew binomial heaps
11 * as described by:
12 * G. Brodal and C. Okasaki , "Optimal Purely Functional Priority Queues",
13 *    Journal of Functional Programming 6:839-857 (1996),
14 *
15 * Based on the heaps Haskell library by Edward Kmett
16 */
17sealed abstract class Heap[A] {
18
19  import Heap._
20  import Heap.impl._
21  import Tree._
22
23  def fold[B](empty: => B, nonempty: (Int, (A, A) => Boolean, Tree[Ranked[A]]) => B): B
24
25  /**Is the heap empty? O(1)*/
26  def isEmpty: Boolean = fold(true, (_, _, _) => false)
27
28  /**Is the heap populated? O(1)*/
29  final def nonEmpty: Boolean = !isEmpty
30
31  /**The number of elements in the heap. O(1)*/
32  def size: Int = fold(0, (s, _, _) => s)
33
34  /**Insert a new value into the heap. O(1)*/
35  def insert(a: A)(implicit o: Order[A]): Heap[A] = insertWith(o.lessThanOrEqual, a)
36
37  /** Alias for insert */
38  final def +(a: A)(implicit o: Order[A]): Heap[A] = this insert a
39
40  def insertAll(as: TraversableOnce[A])(implicit o: Order[A]): Heap[A] =
41    as.foldLeft(this)((h,a) => h insert a)
42
43  def insertAllF[F[_]](as: F[A])(implicit F: Foldable[F], o: Order[A]): Heap[A] =
44    F.foldLeft(as, this)((h,a) => h insert a)
45
46  /**Meld the values from two heaps into one heap. O(1)*/
47  def union(as: Heap[A]): Heap[A] = (this, as) match {
48    case (Empty(), q)                         => q
49    case (q, Empty())                         => q
50    case (Heap(s1, leq, t1@Node(Ranked(r1, x1), f1)),
51    Heap(s2, _, t2@Node(Ranked(r2, x2), f2))) =>
52      if (leq(x1, x2))
53        Heap(s1 + s2, leq, Node(Ranked(0, x1), skewInsert(leq, t2, f1)))
54      else
55        Heap(s1 + s2, leq, Node(Ranked(0, x2), skewInsert(leq, t1, f2)))
56  }
57
58  /**Split the heap into the minimum element and the remainder. O(log n)*/
59  def uncons: Option[(A, Heap[A])] =
60    fold(None, (_, _, t) => Some((t.rootLabel.value, deleteMin)))
61
62  /**Get the minimum key on the (nonempty) heap. O(1) */
63  def minimum: A = fold(sys.error("Heap.minimum: empty heap"), (_, _, t) => t.rootLabel.value)
64
65  /**Get the minimum key on the (nonempty) heap. O(1) */
66  def minimumO: Option[A] = fold(None, (_, _, t) => Some(t.rootLabel.value))
67
68  /**Delete the minimum key from the heap and return the resulting heap. O(log n) */
69  def deleteMin: Heap[A] = {
70    fold(Empty[A], (s, leq, t) => t match {
71      case Leaf(_) => Empty[A]
72      case Node(_, f0)       => {
73        val (Node(Ranked(r, x), cf), ts2) = getMin(leq, f0)
74        val (zs, ts1, f1) : (Forest[A], Forest[A], Forest[A]) = splitForest(r, emptyEphemeralStream, emptyEphemeralStream, cf)
75        val f2 = skewMeld(leq, skewMeld(leq, ts1, ts2), f1)
76        val f3 = zs.foldRight(f2)(skewInsert(leq, _, _))
77        Heap(s - 1, leq, Node(Ranked(0, x), f3))
78      }
79    })
80  }
81
82  def adjustMin(f: A => A): Heap[A] = this match {
83    case Heap(s, leq, Node(Ranked(r, x), xs)) =>
84      Heap(s, leq, heapify(leq)(Node(Ranked(r, f(x)), xs)))
85  }
86
87  def toUnsortedStream: EStream[A] = fold(emptyEphemeralStream, (_, _, t) => t.flatten.map(_.value))
88
89  def toUnsortedList: List[A] = toUnsortedStream.toList
90
91  def toUnsortedIList: IList[A] = IList.fromFoldable(toUnsortedStream)
92
93  def toEphemeralStream: EStream[A] =
94    EphemeralStream.unfold(this)(_.uncons)
95
96  def toStream: Stream[A] =
97    Foldable[EphemeralStream].toStream(this.toEphemeralStream)
98
99  def toList: List[A] = toStream.toList
100
101  def toIList: IList[A] = toEphemeralStream.toIList
102
103  /**Map a function over the heap, returning a new heap ordered appropriately. O(n)*/
104  def map[B: Order](f: A => B): Heap[B] = fold(Empty[B], (_, _, t) => t.foldMap(x => singleton(f(x.value))))
105
106  def forall(f: A => Boolean): Boolean = toStream.forall(f)
107
108  def exists(f: A => Boolean): Boolean = toStream.exists(f)
109
110  def foreach(f: A => Unit): Unit = toStream.foreach(f)
111
112  /**Filter the heap, retaining only values that satisfy the predicate. O(n)*/
113  def filter(p: A => Boolean): Heap[A] =
114    fold(Empty[A], (_, leq, t) => t foldMap (x => if (p(x.value)) singletonWith(leq, x.value) else Empty[A]))
115
116  /**Partition the heap according to a predicate. The first heap contains all elements that
117   * satisfy the predicate. The second contains all elements that fail the predicate. O(n)*/
118  def partition(p: A => Boolean): (Heap[A], Heap[A]) =
119    fold((Empty[A], Empty[A]), (_, leq, t) => t.foldMap(x =>
120      if (p(x.value)) (singletonWith(leq, x.value), Empty[A])
121      else
122        (Empty[A], singletonWith(leq, x.value))))
123
124  /**Partition the heap of the elements that are less than, equal to, and greater than a given value. O(n)*/
125  def split(a: A): (Heap[A], Heap[A], Heap[A]) = {
126    fold((Empty[A], Empty[A], Empty[A]), (s, leq, t) => {
127      def f(x: A) : (Heap[A], Heap[A], Heap[A]) = if (leq(x, a)) if (leq(a, x)) (Empty[A], singletonWith(leq, x), Empty[A])
128      else
129        (singletonWith(leq, x), Empty[A], Empty[A])
130      else
131        (Empty[A], Empty[A], singletonWith(leq, x))
132      t foldMap (x => f(x.value))
133    })
134  }
135
136  /**Return a heap consisting of the least n elements of this heap. O(n log n) */
137  def take(n: Int): Heap[A] = withList(_.take(n))
138
139  /**Return a heap consisting of all the members of this heap except for the least n. O(n log n) */
140  def drop(n: Int): Heap[A] = withList(_.drop(n))
141
142  /**Split into two heaps, the first containing the n least elements, the second containing the n
143   * greatest elements. O(n log n) */
144  def splitAt(n: Int): (Heap[A], Heap[A]) = splitWithList(_.splitAt(n))
145
146  /**Returns a tuple where the first element is a heap consisting of the longest prefix of least elements
147   * in this heap that do not satisfy the given predicate, and the second element is the remainder
148   * of the elements. O(n log n) */
149  def break(p: A => Boolean): (Heap[A], Heap[A]) =
150    span(x => !p(x))
151
152  /**Returns a tuple where the first element is a heap consisting of the longest prefix of least elements
153   * in this heap that satisfy the given predicate and the second element is the remainder of the elements.
154   * O(n log n)*/
155  def span(p: A => Boolean): (Heap[A], Heap[A]) =
156    splitWithList(_.span(p))
157
158  /**Returns a heap consisting of the longest prefix of least elements of this heap that satisfy the predicate.
159   * O(n log n) */
160  def takeWhile(p: A => Boolean): Heap[A] =
161    withList(_.takeWhile(p))
162
163  /**Returns a heap consisting of the longest prefix of least elements of this heap that do not
164   * satisfy the predicate. O(n log n) */
165  def dropWhile(p: A => Boolean): Heap[A] =
166    withList(_.dropWhile(p))
167
168  /**Remove duplicate entries from the heap. O(n log n)*/
169  def nub: Heap[A] = fold(Empty[A], (_, leq, t) => {
170    val x = t.rootLabel.value
171    val xs = deleteMin
172    val zs = xs.dropWhile(leq(_, x))
173    zs.nub.insertWith(leq, x)
174  })
175
176  /**Construct heaps from each element in this heap and union them together into a new heap. O(n)*/
177  def flatMap[B: Order](f: A => Heap[B]): Heap[B] =
178    fold(Empty[B], (_, _, t) => t foldMap (x => f(x.value)))
179
180  /**Traverse the elements of the heap in sorted order and produce a new heap with applicative effects.
181   * O(n log n)*/
182  def traverse[F[_] : Applicative, B: Order](f: A => F[B]): F[Heap[B]] = {
183    val F = Applicative[F]
184    F.map(F.traverse(toEphemeralStream)(f))(fromCodata[EStream, B])
185  }
186
187  def foldRight[B](z: B)(f: (A, => B) => B): B = {
188    Foldable[EStream].foldRight(toEphemeralStream, z)(f)
189  }
190
191  private def withList(f: List[A] => List[A]) = {
192    import std.list._
193    fold(Empty[A], (_, leq, _) => fromDataWith(leq, f(toList)))
194  }
195
196  private[scalaz] def insertWith(f: (A, A) => Boolean, x: A) =
197    fold(singletonWith(f, x), (s, _, t) => {
198      val y = t.rootLabel.value
199      if (f(x, y)) Heap(s + 1, f, Node(Ranked(0, x), (t ##:: emptyEphemeralStream)))
200      else
201        Heap(s + 1, f, Node(Ranked(0, y),
202          skewInsert(f, Node(Ranked(0, x), emptyEphemeralStream[Tree[Ranked[A]]]), t.subForest)))
203    })
204
205  private def splitWithList(f: List[A] => (List[A], List[A])) = {
206    import std.list._
207
208    fold((Empty[A], Empty[A]), (_, leq, _) => {
209      val g = (x: List[A]) => fromDataWith(leq, x)
210      val x = f(toList)
211      (g(x._1), g(x._2))
212    })
213  }
214
215  override def toString = "<heap>"
216}
217
218object Heap extends HeapInstances {
219
220  final case class Ranked[A](rank: Int, value: A)
221
222  type Forest[A] = EStream[Tree[Ranked[A]]]
223  type ForestZipper[A] = (Forest[A], Forest[A])
224
225  /**The empty heap */
226  object Empty {
227    def apply[A]: Heap[A] = new Heap[A] {
228      def fold[B](empty: => B, nonempty: (Int, (A, A) => Boolean, Tree[Ranked[A]]) => B): B = empty
229    }
230
231    def unapply[A](h: Heap[A]): Boolean = h.fold(true, (_, _, _) => false)
232  }
233
234  import Heap.impl._
235
236  def fromData[F[_] : Foldable, A: Order](as: F[A]): Heap[A] =
237    Foldable[F].foldLeft(as, Empty[A])((b, a) => b insert a)
238
239  def fromCodata[F[_] : Foldable, A: Order](as: F[A]): Heap[A] =
240    Foldable[F].foldr(as, Empty[A])(x => y => y insert x)
241
242  def fromDataWith[F[_] : Foldable, A](f: (A, A) => Boolean, as: F[A]): Heap[A] =
243    Foldable[F].foldLeft(as, Empty[A])((x, y) => x.insertWith(f, y))
244
245  /**Heap sort */
246  def sort[F[_] : Foldable, A: Order](xs: F[A]): IList[A] = fromData(xs).toIList
247
248  /**Heap sort */
249  def sortWith[F[_] : Foldable, A](f: (A, A) => Boolean, xs: F[A]): IList[A] = fromDataWith(f, xs).toIList
250
251  /**A heap with one element. */
252  def singleton[A: Order](a: A): Heap[A] = singletonWith[A](Order[A].lessThanOrEqual, a)
253
254  /**Create a heap consisting of multiple copies of the same value. O(log n) */
255  def replicate[A: Order](a: A, i: Int): Heap[A] = {
256    def f(x: Heap[A], y: Int): Heap[A] =
257      if (y % 2 == 0) f(x union x, y / 2)
258      else
259      if (y == 1) x
260      else
261        g(x union x, (y - 1) / 2, x)
262    def g(x: Heap[A], y: Int, z: Heap[A]): Heap[A] =
263      if (y % 2 == 0) g(x union x, y / 2, z)
264      else
265      if (y == 1) x union z
266      else
267        g(x union x, (y - 1) / 2, x union z)
268    if (i < 0) sys.error("Heap.replicate: negative length")
269    else
270    if (i == 0) Empty[A]
271    else
272      f(singleton(a), i)
273  }
274
275  def apply[A](sz: Int, leq: (A, A) => Boolean, t: Tree[Ranked[A]]): Heap[A] = new Heap[A] {
276    def fold[B](empty: => B, nonempty: (Int, (A, A) => Boolean, Tree[Ranked[A]]) => B) =
277      nonempty(sz, leq, t)
278  }
279
280  def unapply[A](h: Heap[A]): Option[(Int, (A, A) => Boolean, Tree[Ranked[A]])] =
281    h.fold(None, (sz, leq, t) => Some((sz, leq, t)))
282
283  private[scalaz] object impl {
284    import Tree._
285
286    def rightZ[A]: ForestZipper[A] => ForestZipper[A] = {
287      case (path, x ##:: xs) => (x ##:: path, xs)
288    }
289
290    def adjustZ[A](f: Tree[Ranked[A]] => Tree[Ranked[A]]):
291    ForestZipper[A] => ForestZipper[A] = {
292      case (path, x ##:: xs) => (path, f(x) ##:: xs)
293      case z                => z
294    }
295
296    def rezip[A]: ForestZipper[A] => Forest[A] = {
297      case (s, xs) if s.isEmpty   => xs
298      case (x ##:: path, xs) => rezip((path, x ##:: xs))
299    }
300
301    def rootZ[A]: ForestZipper[A] => A = {
302      case (_, x ##:: _) => x.rootLabel.value
303      case _            => sys.error("Heap.rootZ: empty zipper")
304    }
305
306    def zipper[A](xs: Forest[A]): ForestZipper[A] = (emptyEphemeralStream, xs)
307
308    def emptyZ[A]: ForestZipper[A] = (emptyEphemeralStream, emptyEphemeralStream)
309
310    def minZ[A](f: (A, A) => Boolean): Forest[A] => ForestZipper[A] = {
311      case s if s.isEmpty => emptyZ
312      case xs       => {
313        val z = zipper(xs)
314        minZp(f)(z, z)
315      }
316    }
317
318    def minZp[A](leq: (A, A) => Boolean):
319    (ForestZipper[A], ForestZipper[A]) => ForestZipper[A] = {
320      case (lo, (_, ts)) if ts.isEmpty => lo
321      case (lo, z)             => minZp(leq)(if (leq(rootZ(lo), rootZ(z))) lo else z, rightZ(z))
322    }
323
324    def heapify[A](leq: (A, A) => Boolean): Tree[Ranked[A]] => Tree[Ranked[A]] = {
325      case n@Node(_, ts) if ts.isEmpty      => n
326      case n@Node(Ranked(r, a), as) => {
327        val (left, Node(Ranked(rp, ap), asp) ##:: right) = minZ(leq)(as)
328        if (leq(a, ap)) n
329        else
330          Node(Ranked(r, ap), rezip((left, heapify(leq)(Node(Ranked(rp, a), asp)) ##:: right)))
331      }
332    }
333
334    def singletonWith[A](f: (A, A) => Boolean, a: A) =
335      Heap(1, f, Node(Ranked(0, a), emptyEphemeralStream[Tree[Ranked[A]]]))
336
337
338    def rank[A](t: Tree[Ranked[A]]) = t.rootLabel.rank
339
340    def skewLink[A](f: (A, A) => Boolean,
341                    t0: Tree[Ranked[A]],
342                    t1: Tree[Ranked[A]],
343                    t2: Tree[Ranked[A]]): Tree[Ranked[A]] = (t0, t1, t2) match {
344      case (Node(Ranked(r0, x0), cf0), Node(Ranked(r1, x1), cf1), Node(Ranked(r2, x2), cf2)) =>
345        if (f(x1, x0) && f(x1, x2)) Node(Ranked(r1 + 1, x1), t0 ##:: t2 ##:: cf1)
346        else
347        if (f(x2, x0) && f(x2, x1)) Node(Ranked(r2 + 1, x2), t0 ##:: t1 ##:: cf2)
348        else
349          Node(Ranked(r1 + 1, x0), t1 ##:: t2 ##:: cf0)
350    }
351
352    def link[A](f: (A, A) => Boolean):
353    (Tree[Ranked[A]], Tree[Ranked[A]]) => Tree[Ranked[A]] = {
354      case (t1@Node(Ranked(r1, x1), cf1), t2@Node(Ranked(r2, x2), cf2)) =>
355        if (f(x1, x2)) Node(Ranked(r1 + 1, x1), t2 ##:: cf1)
356        else
357          Node(Ranked(r2 + 1, x2), t1 ##:: cf2)
358    }
359
360    def skewInsert[A](f: (A, A) => Boolean, t: Tree[Ranked[A]], ts: Forest[A]): Forest[A] =
361      ts match {
362        case t1 ##:: t2 ##:: rest =>
363          if (rank(t1) == rank(t2))
364            skewLink(f, t, t1, t2) ##:: rest
365          else (t ##:: ts)
366        case _                  => t ##:: ts
367      }
368
369    def getMin[A](f: (A, A) => Boolean, trees: Forest[A]): (Tree[Ranked[A]], Forest[A]) =
370      trees match {
371        case (t ##:: ts) if ts.isEmpty => (t, emptyEphemeralStream)
372        case t ##:: ts  => {
373          val (tp, tsp) = getMin(f, ts)
374          if (f(t.rootLabel.value, tp.rootLabel.value)) (t, ts) else (tp, t ##:: tsp)
375        }
376      }
377
378    def splitForest[A]:
379    (Int, Forest[A], Forest[A], Forest[A]) => (Forest[A], Forest[A], Forest[A]) = {
380      case (0, zs, ts, f)                  => (zs, ts, f)
381      case (1, zs, ts, (t ##:: tss)) if tss.isEmpty => (zs, t ##:: ts, emptyEphemeralStream)
382      case (1, zs, ts, t1 ##:: t2 ##:: f)    =>
383        if (rank(t2) == 0) (t1 ##:: zs, t2 ##:: ts, f)
384        else
385          (zs, t1 ##:: ts, t2 ##:: f)
386      case (r, zs, ts, (t1 ##:: t2 ##:: cf)) =>
387        if (rank(t1) == rank(t2)) (zs, t1 ##:: t2 ##:: ts, cf)
388        else
389        if (rank(t1) == 0) splitForest(r - 1, t1 ##:: zs, t2 ##:: ts, cf)
390        else
391          splitForest(r - 1, zs, t1 ##:: ts, t2 ##:: cf)
392      case (_, _, _, _)                    => sys.error("Heap.splitForest: invalid arguments")
393    }
394
395    def skewMeld[A](f: (A, A) => Boolean, ts: Forest[A], tsp: Forest[A]) =
396      unionUniq(f)(uniqify(f)(ts), uniqify(f)(tsp))
397
398    def ins[A](f: (A, A) => Boolean, t: Tree[Ranked[A]]): Forest[A] => Forest[A] = {
399      case s if s.isEmpty    => EphemeralStream(t)
400      case (tp ##:: ts) => if (rank(t) < rank(tp)) t ##:: tp ##:: ts
401      else
403    }
404
405    def uniqify[A](f: (A, A) => Boolean): Forest[A] => Forest[A] = {
406      case s if s.isEmpty   => emptyEphemeralStream
407      case (t ##:: ts) => ins(f, t)(ts)
408    }
409
410    def unionUniq[A](f: (A, A) => Boolean): (Forest[A], Forest[A]) => Forest[A] = {
411      case (s, ts) if s.isEmpty                         => ts
412      case (ts, s) if s.isEmpty                         => ts
413      case (tts1@(t1 ##:: ts1), tts2@(t2 ##:: ts2)) =>
414        import std.anyVal._
415        Order[Int].order(rank(t1), rank(t2)) match {
416          case Ordering.LT => t1 ##:: unionUniq(f)(ts1, tts2)
417          case Ordering.EQ => ins(f, link(f)(t1, t2))(unionUniq(f)(ts1, ts2))
418          case Ordering.GT => t2 ##:: unionUniq(f)(tts1, ts2)
419        }
420    }
421  }
422}
423
424sealed abstract class HeapInstances {
425  implicit val heapInstance: Foldable[Heap] = new Foldable[Heap] with Foldable.FromFoldr[Heap] {
426    def foldRight[A, B](fa: Heap[A], z: => B)(f: (A, => B) => B) = fa.foldRight(z)(f)
427  }
428
429  implicit def heapMonoid[A]: Monoid[Heap[A]] = new Monoid[Heap[A]] {
430    def append(f1: Heap[A], f2: => Heap[A]) = f1 union f2
431    def zero = Heap.Empty.apply
432  }
433
434  // implicit def heapEqual[A: Equal]: Equal[Heap[A]] = Equal.equalBy((_: Heap[A]).toStream)
435  implicit def healEqual[A : Equal](implicit H : Foldable[Heap]): Equal[Heap[A]] = Equal[EStream[A]] contramap {H.toEphemeralStream(_: Heap[A])}
436}
```