From 4e20ede1f2149ec68c08eb90cf4bd0c2c565cb73 Mon Sep 17 00:00:00 2001
From: "J. Ian Johnson" <ianj@ccs.neu.edu>
Date: Mon, 31 Mar 2014 15:47:39 -0400
Subject: [PATCH] Extend heaps to allow removing arbitrary elements via an
 additional equality operator.

---
 .../data-doc/data/scribblings/heap.scrbl      |  9 ++++++
 pkgs/data-pkgs/data-lib/data/heap.rkt         | 30 +++++++++++++++++++
 pkgs/data-pkgs/data-test/tests/data/heap.rkt  |  6 ++++
 3 files changed, 45 insertions(+)

diff --git a/pkgs/data-pkgs/data-doc/data/scribblings/heap.scrbl b/pkgs/data-pkgs/data-doc/data/scribblings/heap.scrbl
index 428dd1a92e..29198df793 100644
--- a/pkgs/data-pkgs/data-doc/data/scribblings/heap.scrbl
+++ b/pkgs/data-pkgs/data-doc/data/scribblings/heap.scrbl
@@ -109,6 +109,15 @@ empty, an exception is raised.
   (heap-min a-heap)]
 }
 
+@defproc[(heap-remove! [h heap?] [v any/c] [#:same? same? (-> any/c any/c any/c) equal?]) void?]{
+Removes @racket[v] from the heap @racket[h] if it exists. 
+@examples[#:eval the-eval
+  (define a-heap (make-heap string<=? string=?))
+  (heap-add! a-heap "a" "b" "c")
+  (heap-remove! a-heap "b")
+  (for/list ([a (in-heap a-heap)]) a)]
+}
+
 @defproc[(vector->heap [<=? (-> any/c any/c any/c)] [items vector?]) heap?]{
 
 Builds a heap with the elements from @racket[items]. The vector is not
diff --git a/pkgs/data-pkgs/data-lib/data/heap.rkt b/pkgs/data-pkgs/data-lib/data/heap.rkt
index d3ebd95655..6f6ee0d025 100644
--- a/pkgs/data-pkgs/data-lib/data/heap.rkt
+++ b/pkgs/data-pkgs/data-lib/data/heap.rkt
@@ -156,6 +156,35 @@
        (set-heap-vec! h (shrink-vector vec)))
      (set-heap-count! h (sub1 size))]))
 
+(define (heap-get-index h v same?)
+  (match h
+    [(heap vec size <=?)
+     (and (not (eq? 0 size))
+          (let search ([n 0] [n-key (vector-ref vec 0)])
+            (cond
+             [(same? n-key v) n]
+             ;; The heap property ensures n-key <= all its children
+             [else
+              (define (search-right)
+                (define right (vt-rightchild n))
+                (and (< right size)
+                     (let ([right-key (vector-ref vec right)])
+                       (and (<=? right-key v)
+                            (search right right-key)))))
+              ;; Try going left if the left child is <= v
+              (define left (vt-leftchild n))
+              (and (< left size) ;; if no left, there can't be a right.
+                   (let ([left-key (vector-ref vec left)])
+                     ;; If left <= v, try left side.
+                     (if (<=? left-key v)
+                         (or (search left left-key) (search-right))
+                         (search-right))))])))]))
+
+(define (heap-remove! h v #:same? [same? equal?])
+  (match (heap-get-index h v same?)
+    [#f (void)]
+    [n (heap-remove-index! h n)]))
+
 (define (in-heap h)
   (in-heap/consume! (heap-copy h)))
 
@@ -211,6 +240,7 @@
  [heap-add-all! (-> heap? (or/c list? vector? heap?) void?)]
  [heap-min (-> heap? any/c)]
  [heap-remove-min! (-> heap? void?)]
+ [heap-remove! (->* (heap? any/c) [#:same? (-> any/c any/c any/c)] void?)]
 
  [vector->heap (-> (-> any/c any/c any/c) vector? heap?)]
  [heap->vector (-> heap? vector?)]
diff --git a/pkgs/data-pkgs/data-test/tests/data/heap.rkt b/pkgs/data-pkgs/data-test/tests/data/heap.rkt
index 9aace3aa74..577c9ce4a1 100644
--- a/pkgs/data-pkgs/data-test/tests/data/heap.rkt
+++ b/pkgs/data-pkgs/data-test/tests/data/heap.rkt
@@ -32,6 +32,12 @@
     (heap->vector h))
   '#(4 6 8 10))
 
+(test-equal? "heap-remove!"
+  (let ([h (mkheap)])
+    (heap-remove! h 4)
+    (heap->vector h))
+  '#(2 6 8 10))
+
 (define (rand-test range count1 count2 count3)
   (let ([h (make-heap <=)]
         [xs null]) ;; mutated