Today’s Programming Praxis problem is about Treaps – binary trees that are more or less balanced thanks to random numbers. The provided solution is 58 lines, so we have quite a bit of work to do. Let’s get started.
First our imports:
import Control.Monad import Data.Char import qualified Data.List.Key as K import System.Random
As mentioned in the intro, the data structure is, for all intents and purposes, a binary tree.
data Treap k a = Nil | Node Int k a (Treap k a) (Treap k a)
Nil nodes have a priority of -1.
priority :: Treap k a -> Int priority Nil = -1 priority (Node p _ _ _ _) = p
Rotating left and right requires some node reshuffling:
rotLeft :: Treap k a -> Treap k a rotLeft (Node p k a l (Node rp rk ra rl rr)) = Node rp rk ra (Node p k a l rl) rr rotLeft t = t rotRight :: Treap k a -> Treap k a rotRight (Node p k a (Node lp lk la ll lr) r) = Node lp lk la ll (Node p k a lr r) rotRight t = t
This function automatically applies the correct rotations:
rot :: Treap k a -> Treap k a rot Nil = Nil rot t@(Node p _ _ l r) | p < priority l = rotRight t | p < priority r = rotLeft t | otherwise = t
Lookup works the same as in any binary tree.
find :: Ord k => k -> Treap k a -> Maybe a find _ Nil = Nothing find k' (Node _ k a l r) | k' < k = find k' l | k' > k = find k' r | otherwise = Just a
And updating is also not that different. The only annoyance is that we have to do it in the IO monad, since we need random numbers.
update :: Ord k => (a -> a -> a) -> k -> a -> Treap k a -> IO (Treap k a) update _ k' a' Nil = fmap (\r -> Node r k' a' Nil Nil) $ randomRIO (0, maxBound) update f k' a' (Node p k a l r) | k' < k = fmap (\n -> rot $ Node p k a n r) (update f k' a' l) | k' > k = fmap (rot . Node p k a l) (update f k' a' r) | otherwise = return $ Node p k' (f a' a) l r
Inserting is just a special case of updating:
insert :: Ord k => k -> a -> Treap k a -> IO (Treap k a) insert = update const
For deleting, we need a way the combine the two remaining subtrees. This requires some rotation.
deroot :: Treap k a -> Treap k a deroot Nil = Nil deroot t@(Node _ _ _ l r) | priority l < priority r = d deroot id $ rotLeft t | otherwise = d id deroot $ rotRight t where d fl fr = (\(Node p k a l' r') -> Node p k a (fl l') (fr r'))
That makes deleting a simple matter of applying deroot at the correct location.
delete :: Ord k => k -> Treap k a -> Treap k a delete _ Nil = Nil delete k' t@(Node p k a l r) | k' < k = Node p k a (delete k' l) r | k' > k = Node p k a l (delete k' r) | otherwise = deroot t
Converting a treap to a list is just simple recursion.
toList :: Treap k a -> [(k, a)] toList Nil =  toList (Node _ k a l r) = toList l ++ [(k, a)] ++ toList r
And there we go. Now to test it on the word frequency problem:
main :: IO () main = mapM_ print =<< wordFreqs 25 =<< readFile "bible.txt" wordFreqs :: Int -> String -> IO [(String, Int)] wordFreqs n = fmap (take n . reverse . K.sort snd . toList) . foldM (\a w -> update (+) w 1 a) Nil . map (filter isAlpha) . words
Works as expected. At 40 lines it’s about a 30% reduction compared to the scheme code. Not bad.