## Posts Tagged ‘treaps’

### Programming Praxis – Treaps

June 26, 2009

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.