Files
newt/playground/samples/aoc2024/SortedMap.newt

68 lines
2.6 KiB
Agda
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
module SortedMap
import Prelude
data T23 : Nat -> U -> U -> U where
Leaf : k v. k -> v -> T23 Z k v
Node2 : h k v. T23 h k v -> k -> T23 h k v -> T23 (S h) k v
Node3 : h k v. T23 h k v -> k -> T23 h k v -> k -> T23 h k v -> T23 (S h) k v
lookupT23 : h k v. {{Ord k}} {{Eq k}} -> k -> T23 h k v -> Maybe (k × v)
lookupT23 key (Leaf k v)= if k == key then Just (k,v) else Nothing
lookupT23 key (Node2 t1 k1 t2) =
if key <= k1 then lookupT23 key t1 else lookupT23 key t2
lookupT23 key (Node3 t1 k1 t2 k2 t3) =
if key <= k1 then lookupT23 key t1
else if key <= k2 then lookupT23 key t2
else lookupT23 key t3
insertT23 : h k v. {{Ord k}} {{Eq k}} -> k -> v -> T23 h k v -> Either (T23 h k v) (T23 h k v × k × T23 h k v)
insertT23 key value (Leaf k v) =
if key == k then Left (Leaf key value)
else if key <= k then Right (Leaf key value, key, Leaf k v)
else Right (Leaf k v, k, Leaf key value)
insertT23 key value (Node2 t1 k1 t2) =
if key <= k1 then
case insertT23 key value t1 of
Left t1' => Left (Node2 t1' k1 t2)
Right (a,b,c) => Left (Node3 a b c k1 t2)
else case insertT23 key value t2 of
Left t2' => Left (Node2 t1 k1 t2')
Right (a,b,c) => Left (Node3 t1 k1 a b c)
insertT23 key value (Node3 t1 k1 t2 k2 t3) =
if key <= k1 then
case insertT23 key value t1 of
Left t1' => Left (Node3 t1' k1 t2 k2 t3)
Right (a,b,c) => Right (Node2 a b c, k1, Node2 t2 k2 t3)
else if key <= k2 then
case insertT23 key value t2 of
Left t2' => Left (Node3 t1 k1 t2' k2 t3)
Right (a,b,c) => Right (Node2 t1 k1 a, b, Node2 c k2 t3)
else
case insertT23 key value t3 of
Left t3' => Left (Node3 t1 k1 t2 k2 t3')
Right (a,b,c) => Right (Node2 t1 k1 t2, k2, Node2 a b c)
data SortedMap : U -> U -> U where
EmptyMap : k v. SortedMap k v
MapOf : k v h. T23 h k v -> SortedMap k v
lookupMap : k v. {{Ord k}} {{Eq k}} -> k -> SortedMap k v -> Maybe (k × v)
lookupMap k EmptyMap = Nothing
lookupMap k (MapOf map) = lookupT23 k map
updateMap : k v. {{Ord k}} {{Eq k}} -> k -> v -> SortedMap k v -> SortedMap k v
updateMap k v EmptyMap = MapOf $ Leaf k v
updateMap k v (MapOf map) = case insertT23 k v map of
Left map' => MapOf map'
Right (a, b, c) => MapOf (Node2 a b c)
toList : k v. SortedMap k v List (k × v)
toList {k} {v} (MapOf smap) = reverse $ go smap Nil
where
go : h. T23 h k v List (k × v) List (k × v)
go (Leaf k v) acc = (k, v) :: acc
go (Node2 t1 k1 t2) acc = go t2 (go t1 acc)
go (Node3 t1 k1 t2 k2 t3) acc = go t3 $ go t2 $ go t1 acc
toList _ = Nil