okasaki

Okasaki's Purely Functional Data Structures
git clone git://git.jtobin.io/okasaki.git
Log | Files | Refs | LICENSE

Tree.hs (2634B)


      1 {-# OPTIONS_GHC -Wall #-}
      2 {-# LANGUAGE LambdaCase #-}
      3 
      4 module Tree where
      5 
      6 import Data.Maybe
      7 import Prelude hiding (lookup)
      8 
      9 data Tree a = Leaf | Node (Tree a) a (Tree a) deriving (Eq, Show)
     10 
     11 node :: a -> Tree a -> Tree a -> Tree a
     12 node x l = Node l x
     13 
     14 leaf :: Tree a
     15 leaf = Leaf
     16 
     17 insert :: Ord a => a -> Tree a -> Tree a
     18 insert x tree = case tree of
     19   Leaf       -> node x leaf leaf
     20   Node l e r -> case compare x e of
     21     EQ -> tree
     22     LT -> node e (insert x l) r
     23     GT -> node e l (insert x r)
     24 
     25 member :: Ord a => a -> Tree a -> Bool
     26 member x = \case
     27   Leaf       -> False
     28   Node l e r -> case compare x e of
     29     EQ -> True
     30     LT -> member x l
     31     GT -> member x r
     32 
     33 fromList :: Ord a => [a] -> Tree a
     34 fromList = foldr insert leaf
     35 
     36 test :: Tree Int
     37 test = fromList [5..10]
     38 
     39 -- exercise 2.2 (max d + 1 comparisons)
     40 altMember :: Ord a => a -> Tree a -> Bool
     41 altMember x t = go t Nothing where
     42   go (Node l e r) acc
     43     | x < e     = go l acc
     44     | otherwise = go r (Just e)
     45   go Leaf Nothing  = False
     46   go Leaf (Just e) = e == x
     47 
     48 -- exercise 2.3 (no unnecessary copying)
     49 altInsert :: Ord a => a -> Tree a -> Tree a
     50 altInsert x t = fromMaybe t (go t) where
     51   go Leaf = return (Node Leaf x Leaf)
     52   go (Node l e r) = case compare x e of
     53     EQ -> Nothing
     54     LT -> fmap (\alt -> Node alt e r) (go l)
     55     GT -> fmap (\alt -> Node l e alt) (go r)
     56 
     57 -- exercise 2.4 (no unnecessary copying, max d + 1 comparisons)
     58 efficientInsert :: Ord a => a -> Tree a -> Tree a
     59 efficientInsert x t = fromMaybe t (go t Nothing) where
     60   go (Node l e r) acc
     61     | x < e     = fmap (\alt -> Node alt e r) (go l acc)
     62     | otherwise = fmap (\alt -> Node l e alt) (go r (Just e))
     63   go Leaf (Just e)
     64     | e == x    = Nothing
     65     | otherwise = go Leaf Nothing
     66   go Leaf Nothing = return (Node Leaf x Leaf)
     67 
     68 -- exercise 2.5a (balanced binary trees of depth n)
     69 completeDepth :: Ord a => a -> Int -> Tree a
     70 completeDepth x n
     71   | n <= 0 = Leaf
     72   | otherwise =
     73       let t = completeDepth x (pred n)
     74       in  Node t x t
     75 
     76 -- exercise 2.5b (mostly-balanced binary trees of size n)
     77 completeSize :: Ord a => a -> Int -> Tree a
     78 completeSize x n
     79   | n <= 0 = Leaf
     80   | odd n  =
     81       let t = completeSize x (n `quot` 2)
     82       in  Node t x t
     83   | otherwise =
     84       let l = completeSize x (n `quot` 2)
     85           r = completeSize x (n `quot` 2 - 1)
     86       in  Node l x r
     87 
     88 test0 :: Num a => Tree a
     89 test0 = node 1 leaf leaf
     90 
     91 test1 :: Num a => Tree a
     92 test1 = node 2 (node 1 leaf leaf) (node 3 leaf leaf)
     93 
     94 test2 :: (Ord a, Num a) => Tree a
     95 test2 = insert 0 test1
     96 
     97 test3 :: (Ord a, Num a) => Tree a
     98 test3 = insert 5 test1
     99 
    100 test4 :: (Ord a, Num a) => Tree a
    101 test4 = insert 4 test3
    102