Replacing an element in a n-ary tree

Question

I'm pretty new to Haskell and still have some problems getting my head around functional programming. With that said:

I have a custom n-ary tree datatype

data Tree = Empty | Leaf String | Node String [Tree]

I'm trying to write a function to replace an element in a tree, i.e.

replaceInTree :: String -> String -> Tree -> Maybe Tree

Replacing the first string with the second. There is only ever one occurance of each string so I can stick with the first one found. I've made a few efforts but I can't grasp how to reconstruct the full tree after replacing the element. Trivially, I have this:

ntreeReplace x y (Node z lis)
    |(x==z) = Just (Node y lis)

which only changes the head node, obviously. I've written a function that returns true if an element is present in the tree, as a leaf or node, but progress beyond that is proving difficult.

---

Thanks for any help!

Answer 1

This is tricky. You'd like the process to short-circuit on the children of a node if any child produces a match. Here's my solution.

import Data.Maybe

ntreeReplace :: String -> String -> Tree -> Maybe Tree
ntreeReplace x y (Node z lis)
    | (x==z) = Just (Node y lis)
    | otherwise = 
        let (nothings, justs) = span (isNothing . ntreeReplace x y) lis
        in case justs of
             [] -> Nothing
             (t:ts) -> Just (Node z (nothings ++ [fromJust $ ntreeReplace x y t] ++ ts))

ntreeReplace x y (Leaf z)
    | (x==z) = Just (Leaf y)
    | otherwise = Nothing

nTreeReplace returns Nothing if there was no match in this tree (i.e., we should re-use the input unchanged) and Just t if a replacement was made (i.e., we should replace the input with t). I use span to split the children list into a prefix of Nothings and a (possibly empty) suffix where the first element has a match.

This implementation has a possible small inefficiency in that it calls ntreeReplace twice on a matching child: once in the span predicate and again while building the replacement node.

I'd also recommend a higher-level function replace that gives you back a (possibly identical) Tree, instead of a Maybe Tree.

replace :: String -> String -> Tree -> Tree
replace x y t =
    case ntreeReplace x y t of
      Nothing -> t
      (Just t') -> t'

[EDIT] Along the lines of @codebliss' suggestion, you could change the data declaration to

data Tree a = Empty | Leaf a | Node a [Tree a]

The only other thing you would have to change is the signatures of ntreeReplace and replace:

replace :: Eq a => a -> a -> Tree a -> Tree a
ntreeReplace :: Eq a => a -> a -> Tree a -> Maybe (Tree a)

Answer 2

Here's another solution that avoids the double call to ntreeReplace, at the cost of building a bunch of unnecessary list copies. (I have no idea which is more efficient.)

I'm using the modified data definition I suggested above.

import Data.List

data Tree a = Empty | Leaf a | Node a [Tree a]

-- Returns a tree and a boolean flag indicating whether the tree changed
ntreeReplace :: Eq a => a -> a -> Tree a -> (Bool, Tree a)

ntreeReplace x y t@(Node z ts)
    | (x==z) = (True, Node y ts)
    | otherwise = 
        let (changed,ts') =
                foldl'
                (\(changed,ts') t -> 
                     if changed then (True,t:ts')
                     else
                         let (changed,t') = ntreeReplace x y t
                         in (changed,t':ts'))
                (False,[])
                ts
        in 
          if changed then (True, Node z $ reverse ts')
          else (False,t)

ntreeReplace x y t@(Leaf z)
    | (x==z) = (True, Leaf y)
    | otherwise = (False,t)