CODEWITHSUNDEEP
haskelltreetree-traversal
Tobias Hermann
Tobias Hermann

Reputation: 10926

Change leaf values of a tree to ordered sequence while preserving the tree's structure

Given is a recursive tree structure

data Tree = Leaf Int | Node Tree Tree deriving Show

I would like to normalize it in a way that preserves the tree's structure, but makes the integers at the leaves sequential in depths-first order. How can I achieve this? My current setup code looks as follows:

myTree = Node (Leaf 3) (Node (Leaf 5) (Leaf 2))

myTree' = normalize myTree

-- preserve tree structure, but make Ints sequential in depths-first traversal
normalize :: Tree -> Tree
normalize = id -- todo: implement

main = do
    print myTree  -- prints      : Node (Leaf 3) (Node (Leaf 5) (Leaf 2))
    print myTree' -- should print: Node (Leaf 1) (Node (Leaf 2) (Leaf 3))

Upvotes: 1

Views: 340

Answers (2)

András Kovács
András Kovács

Reputation: 30103

The standard way to do this is to record the current label in a state monad:

import Control.Monad.State

data Tree = Leaf Int | Node Tree Tree deriving (Show)

normalize :: Tree -> Tree
normalize t = evalState (go t) 1 where

  go :: Tree -> State Int Tree
  go (Leaf _)   = Leaf <$> (get <* modify (+1))
  go (Node l r) = Node <$> go l <*> go r

This solution initializes the state to 1, then traverses the tree and on each encountered Leaf it puts the current label in the returned Leaf and increments the label.

Alternatively, we can derive the Traversable instance for Tree:

{-# language DeriveFunctor, DeriveFoldable, DeriveTraversable #-}

import Control.Monad.State

data Tree a = Leaf a | Node (Tree a) (Tree a)
  deriving (Show, Functor, Foldable, Traversable)

normalize :: Tree a -> Tree Int
normalize = (`evalState` 1) . traverse (\_ -> get <* modify (+1))

This works the same way, but that note that it relies on the fact that the derived Traversable instance has the right traversal order. Of course, if we wanted some other order then we would have to write our own traversals.

Upvotes: 1

josejuan
josejuan

Reputation: 9566

Without using monads, you can write one function mapping trees carrying certain state

mapLRS :: (s -> Int -> (s, Int)) -> s -> Tree -> (Tree, s)
mapLRS f s (Leaf x  ) =
    let (s', x') = f s x
    in  (Leaf x', s')                  -- the mapped leaf and the new state
mapLRS f s (Node l r) =
    let (l', s' ) = mapLRS f s  l      -- the mapped left branch and the new state
        (r', s'') = mapLRS f s' r      -- the mapped right branch and the new state
    in  (Node l' r', s'')

now

normalize :: Tree -> Tree
normalize = fst . mapWithStateLeftRight (\s _ -> (s + 1, s)) 1

Using monads could be more readable (f is not the same for simplicity using state)

import Control.Monad.State

mapLRS :: (s -> (Int, s)) -> Tree -> State s Tree
mapLRS f (Leaf   x) = Leaf <$> state f
mapLRS f (Node l r) = Node <$> mapLRS f l <*> mapLRS f r

normalize :: Tree -> Tree
normalize tree = evalState (mapLRS (\s -> (s, s + 1)) tree) 1

Upvotes: 1

Related Questions