Haskell: Turning a tree into a map

Question

Basically I want to turn a BST tree into a map where the nodes are the keys and the number of occurrences of the nodes are the values. So if I inputted this:

toMap (leaf 13)

I would get

> [(13,1)]

Here is what I have so far:

data Tree a = Empty | Node a (Tree a) (Tree a) deriving (Show)
leaf x = Node x Empty Empty

toMap' :: Int -> Tree a -> ([(a, Int)], Int)
toMap' a Empty = ([], a)
toMap' a (Node x xl xr) = ((x, a): xl' ++ xr', k)
                      where (xl', i) = toMap' (a+1) xl
                            (xr', k) = toMap' (i) xr

toMap :: Tree a -> [(a, Int)]
toMap = fst. toMap' 1

This program returns a map but the values are incorrect. Each value is one greater than the previous value (so if there are 3 nodes than the value of the third node will be three). I think I have to place a counter on each new key, but I'm not sure how. Thanks in advance

Answer 1

Assuming you have a function foldt that folds across trees (in some order irrelevant to your current application), and some function insertIncr that inserts or increments the value of a key in a Map a Int, you could apply one to the other.

You would deal with the following type signatures:

import Data.Map

foldt :: (a -> b -> b) -> b -> Tree a -> b
foldt f acc Empty = acc
foldt f acc (Node x l r) = acc'
    where accl = foldt f acc l
          accr = foldt f accl r
          acc' = f x accr

-- insert 1 if not present, increment if present
insertIncr :: Ord a => a -> Map a Int -> Map a Int
insertIncr = undefined

toMap' :: Ord a => Tree a -> Map a Int
toMap' = foldt insertIncr empty

The insertIncr function could be made using e.g. Data.Map.insertWith. Notice that the Ord type class is necessary in order to insert something into a Data.Map. If you prefer the plain [(a,Int)] kind of map, then insertIncr could have the type Eq a => a -> [(a,Int)] -> [(a,Int)].

Edit: Changed suggestion of using adjustWithKey to insertWith.

Answer 2

Honestly, this is a case that I'd just solve with a compositional breakdown.

data Tree a = Empty | Node a (Tree a) (Tree a) deriving (Show)

toMap :: Ord a => Tree a -> [(a, Int)]
toMap = countDups . toList

Note that I had to add an extra constraint on a. It needs at least Eq to be solvable at all, but Ord allows asymptotically better solutions.

The basic idea here is just breaking the solution down into parts, and figuring out each part later. So, the next part is toList. I'm not going to assume the order matters, and as such I'll choose a prefix ordering, since it's easy to make both lazy and simple.

toList :: Tree a -> [a]
toList Empty = []
toList (Node a l r) = a : toList l ++ toList r

Ok, nice and straightforward. On to counting the duplicates. Let's break this down into a few pieces, too.

countDups :: Ord a => [a] -> [(a, Int)]
countDups = map (\xs -> (head xs, length xs)) . group . sort

Ok, I may have slightly cheated there by taking advantage of group and sort from Data.List. But on the other hand, this is exactly the kind of problem group is meant to solve. And sorting is just a standard tool for everything.

If I had Control.Arrow imported, I'd replace the lambda with (head &&& length). But that's just a bit of a standard idiom that doesn't really simplify things - it just makes them a bit more succinct to type.

The main idea of this approach is breaking down the problem into pieces that do some meaningful thing on their own. Then compose those pieces together into a full solution. It's handy to have a way to convert a Tree a into a [a]. Might as well have a function that does that. And once you do, the remaining piece is a useful bit of logic to have available for processing lists. And if you break that down, you find it's an easy composition of existing bits list functionality.

This is often one of the best ways to solve problems in any programming language - break the big task down into smaller tasks. What's nice about doing it in Haskell is that composing the smaller tasks into the full process is a nicely succinct process.