0 | module Structures.PriorityQueue.SkewBinomialQueue
  1 |
  2 | import Data.Nat
  3 | import Data.Fin
  4 | import Data.Vect
  5 | import Data.Either
  6 | import Decidable.Equality
  7 |
  8 | import Structures.PriorityQueue
  9 |
 10 | %default total
 11 |
 12 | ||| Skew binomial tree per Optimal purely functional priority queues
 13 | ||| (Brodal, G. S., & Okasaki, C. (1996).)
 14 | |||
 15 | ||| Holds a much more complicated number of elements than normal binomial trees,
 16 | ||| but has properties that enable us to implement inserts as a single link
 17 | data SkewBinomialTree : (e : Type) -> (rank : Nat) -> Type where
 18 |   Singleton : (element : e)
 19 |     -> SkewBinomialTree e 0
 20 |   SimpleLink : (root : SkewBinomialTree e rank) -> (leftmost_child : SkewBinomialTree e rank)
 21 |     -> SkewBinomialTree e (S rank)
 22 |   SkewLinkA : (root : SkewBinomialTree e 0) -> (lefter, left : SkewBinomialTree e rank)
 23 |     -> SkewBinomialTree e (S rank)
 24 |   SkewLinkB : (root : SkewBinomialTree e rank) ->
 25 |     (lefter : SkewBinomialTree e 0) -> (left : SkewBinomialTree e rank)
 26 |     -> SkewBinomialTree e (S rank)
 27 | %name SkewBinomialTree bt, ct, dt
 28 |
 29 | SkewBinomialPair : Type -> Type
 30 | SkewBinomialPair e = (rank : Nat ** SkewBinomialTree e rank)
 31 |
 32 | ||| Get the root of a binomial tree
 33 | root : (tree : SkewBinomialTree e n) -> e
 34 | root (Singleton element) = element
 35 | root (SimpleLink bt _) = root bt
 36 | root (SkewLinkA bt _ _) = root bt
 37 | root (SkewLinkB bt _ _) = root bt
 38 |
 39 | ||| SimpleLink two SkewBiomialTrees together, preserving heap ordering
 40 | |||
 41 | ||| Makes the tree with the larger root the child of the other
 42 | simpleLink : Ord e => (x, y : SkewBinomialTree e n) -> SkewBinomialTree e (S n)
 43 | simpleLink x y =
 44 |   let (x_val, y_val) = (root x, root y)
 45 |   in if x_val < y_val
 46 |     then SimpleLink x y
 47 |     else SimpleLink y x
 48 |
 49 | ||| SkewLink two SkewBiomialTrees together, preserving heap ordering
 50 | |||
 51 | ||| Makes the tree with the larger root the child of the other
 52 | skewLink : Ord e => (zero : SkewBinomialTree e 0) -> (x, y : SkewBinomialTree e rank)
 53 |   -> SkewBinomialTree e (S rank)
 54 | skewLink zero x y =
 55 |   let (z_root, x_root, y_root) = (root zero, root x, root y)
 56 |       min_s_root = min x_root y_root
 57 |       min_root = min z_root min_s_root
 58 |   in if min_root == z_root
 59 |     then SkewLinkA zero x y
 60 |     else if min_root == x_root
 61 |       then SkewLinkB x zero y
 62 |       else SkewLinkB y zero x
 63 |
 64 | ||| Get the rank of a SkewBinomialTree
 65 | rank : (tree : SkewBinomialTree e n) -> Nat
 66 | rank (Singleton _) = 0
 67 | rank (SimpleLink root _) = S (rank root)
 68 | rank (SkewLinkA _ right _) = S (rank right)
 69 | rank (SkewLinkB root _ _) = S (rank root)
 70 |
 71 | rankEq : (tree : SkewBinomialTree e n) -> n = rank tree
 72 | rankEq (Singleton element) = Refl
 73 | rankEq (SimpleLink root leftmost_child) =
 74 |   let rest = rankEq root
 75 |   in eqSucc _ _ rest
 76 | rankEq (SkewLinkA root right left) =
 77 |   let rest = rankEq right
 78 |   in eqSucc _ _ rest
 79 | rankEq (SkewLinkB root lefter left) =
 80 |   let rest = rankEq root
 81 |   in eqSucc _ _ rest
 82 |
 83 | toPair : (tree : SkewBinomialTree e n) -> SkewBinomialPair e
 84 | toPair tree =
 85 |   (rank tree **
 86 |     rewrite sym $ rankEq tree in tree)
 87 |
 88 | ||| Break the links of a tree to extract the root
 89 | extractRoot : (tree : SkewBinomialTree e n) -> (e, List (SkewBinomialPair e))
 90 | extractRoot (Singleton element) = (element, [])
 91 | extractRoot (SimpleLink root leftmost_child) =
 92 |   let (element, rest) = extractRoot root
 93 |   in (element, toPair leftmost_child :: rest)
 94 | extractRoot (SkewLinkA root lefter left) =
 95 |   let (element, rest) = extractRoot root
 96 |   in (element, toPair lefter :: toPair left :: rest)
 97 | extractRoot (SkewLinkB root lefter left) =
 98 |   let (element, rest) = extractRoot root
 99 |   in (element, toPair lefter :: toPair left :: rest)
100 |
101 | ||| Return the number of elements in a skew binomial tree
102 | skewLen : {n : _} -> (tree : SkewBinomialTree e n) -> Nat
103 | skewLen tree = skewLen' tree 0 []
104 |   where
105 |   skewLen' : {n : _} -> (tree : SkewBinomialTree e n) -> (acc : Nat) ->
106 |     (stack : List (m : Nat ** SkewBinomialTree e m))
107 |     -> Nat
108 |   skewLen' (Singleton element) acc [] = 
109 |     1 + acc
110 |   skewLen' (Singleton element) acc ((m ** next) :: xs) =
111 |     skewLen' next (1 + acc) xs
112 |   skewLen' (SimpleLink root leftmost_child) acc stack =
113 |     -- TODO: Figure out a way to remove this assert_smaller
114 |     skewLen' root acc (assert_smaller stack ((_ ** leftmost_child) :: stack))
115 |   skewLen' (SkewLinkA root lefter left) acc stack =
116 |     skewLen' root acc
117 |       (assert_smaller stack $
118 |         (_ ** lefter) :: (_ ** left) :: stack)
119 |   skewLen' (SkewLinkB root lefter left) acc stack =
120 |     skewLen' root acc
121 |       (assert_smaller stack $
122 |         (_ ** lefter) :: (_ ** left) :: stack)
123 |
124 | ||| A binomial queue per Optimal purely functional priority queues
125 | ||| (Brodal, G. S., & Okasaki, C. (1996).)
126 | |||
127 | ||| Slightly more complicated than a `BinomialQueue`, but supports inserts in
128 | ||| constant time
129 | export
130 | record SkewBinomialQueue (e : Type) where
131 |   constructor MkSBQ
132 |   o : Ord e
133 |   forest_size : Nat
134 |   forest : Vect forest_size (SkewBinomialPair e)
135 | %name SkewBinomialQueue sq, tq, uq
136 |
137 | ||| Insert a tree into its proper location in a list of binomial pairs
138 | insertTree :
139 |   Ord e => {n : Nat} ->
140 |   (x : SkewBinomialPair e) ->
141 |   (xs : Vect n (SkewBinomialPair e))
142 |   -> (m : Nat ** Vect m (SkewBinomialPair e))
143 | insertTree x [] = (1 ** [x])
144 | insertTree (x_n ** x) ((y_n ** y) :: xs) =
145 |   case decEq x_n y_n of
146 |     Yes Refl => insertTree (S x_n ** simpleLink x y) xs
147 |     No _ =>
148 |       if x_n < y_n
149 |         then (_ ** (x_n ** x) :: (y_n ** y) :: xs)
150 |         else
151 |           let (rest_len ** rest) = insertTree (x_n **x) xs
152 |           in (_ ** (y_n ** y) :: rest)
153 |
154 | ||| meld, phrased in terms of the inner list
155 | meld' :
156 |   Ord e => {n, m : Nat} ->
157 |   (xs : Vect n (SkewBinomialPair e)) ->
158 |   (ys : Vect m (SkewBinomialPair e))
159 |   -> (o : Nat ** Vect o (SkewBinomialPair e))
160 | meld' xs [] = (length xs ** rewrite lengthCorrect xs in xs)
161 | meld' xs (y :: ys) =
162 |   let (_ ** res) = insertTree y xs
163 |   in meld' res ys
164 |
165 | ||| Utility function, seperates out elments from queues
166 | processTree : SkewBinomialPair e -> Either e (SkewBinomialPair e)
167 | processTree (0 ** (Singleton element)) = Left element
168 | processTree pair@((S k) ** tree) = Right pair
169 |
170 | export
171 | PriorityQueue SkewBinomialQueue where
172 |   empty = MkSBQ o 0 []
173 |
174 |   singleton element = MkSBQ o 1 [(0 ** Singleton element)]
175 |
176 |   isNull bq = null bq.forest
177 |
178 |   length (MkSBQ o forest_size forest) =
179 |     let lens = map (\(_ ** x) => skewLen x) forest
180 |     in sum lens
181 |
182 |   findMin (MkSBQ o forest_size forest) =
183 |     case forest of
184 |       [] => Nothing
185 |       ((_ ** x) :: xs) =>
186 |         let x = root x
187 |             xs = map (\(_ ** y) => root y) xs
188 |         in Just $ foldl min x xs
189 |
190 |   meld x y =
191 |     let (new_size ** new_forest) = meld' @{x.o} x.forest y.forest
192 |     in MkSBQ x.o new_size new_forest
193 |
194 |   insert element (MkSBQ o forest_size forest) =
195 |     case forest of
196 |       [] => MkSBQ o 1 [(0 ** Singleton element)]
197 |       [x] => MkSBQ o 2 [(0 ** Singleton element), x]
198 |       (x_rank ** x) :: (y_rank ** y) :: xs =>
199 |         case decEq x_rank y_rank of
200 |           Yes Refl =>
201 |             let rest = skewLink (Singleton element) x y
202 |             in MkSBQ o _ $ (S y_rank ** rest) :: xs
203 |           No _ => MkSBQ o _ $ (0 ** Singleton element) :: forest
204 |
205 |   deleteMin (MkSBQ o forest_size forest) with (forest)
206 |     deleteMin (MkSBQ o 0 forest) | ([]) =
207 |       MkSBQ o 0 []
208 |     deleteMin (MkSBQ o (S 0) forest) | ((x_rank ** x) :: []) =
209 |       let (singlets, rest) =
210 |             partitionEithers . map processTree . reverse . snd . extractRoot
211 |             $ x
212 |           new_queue = MkSBQ o (length rest) (fromList rest)
213 |       in foldl (\acc, e => insert e acc) new_queue singlets
214 |     deleteMin (MkSBQ o (S (S len)) forest) | (x :: y :: xs) =
215 |       let (idx, (min_n ** min_tree)) = findMin (y :: xs) (0, x) 1
216 |           without_min = deleteAt idx forest
217 |           (singlets, rest) =
218 |             partitionEithers . map processTree . reverse . snd . extractRoot
219 |             $ min_tree
220 |           (rest_len ** rest) = listToVect rest
221 |           (new_size ** new_forest) = meld' without_min rest
222 |           new_queue = MkSBQ o new_size new_forest
223 |       in foldl (\acc, e => insert e acc) new_queue singlets
224 |     where
225 |       listToVect : List a -> (n : Nat ** Vect n a)
226 |       listToVect xs = (length xs ** fromList xs)
227 |       findMin :
228 |         {m : Nat} ->
229 |         (xs : Vect n (SkewBinomialPair e)) ->
230 |         (acc : (Fin m, SkewBinomialPair e)) ->
231 |         (idx : Fin m)
232 |         -> (Fin m, SkewBinomialPair e)
233 |       findMin [] acc idx = acc
234 |       findMin ((x_rank ** x) :: xs) orig_acc@(acc_idx, (acc_rank ** acc)) idx =
235 |         if root x < root acc
236 |           then findMin xs (idx, (x_rank ** x)) (finS idx)
237 |           else findMin xs orig_acc (finS idx)
238 |