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.