Breaking after finding the kth element of an inorder traversal using a higher order traversal function

Question

I have the following code to do an inorder traversal of a Binary Tree:

data BinaryTree a =
  Node a (BinaryTree a) (BinaryTree a)
  | Leaf
  deriving (Show)

inorder :: (a -> b -> b) -> b -> BinaryTree a -> b
inorder f acc tree = go tree acc
  where go Leaf z = z
        go (Node v l r) z = (go r . f v . go l) z

Using the inorder function above I'd like to get the kth element without having to traverse the entire list.

The traversal is a little like a fold given that you pass it a function and a starting value. I was thinking that I could solve it by passing k as the starting value, and a function that'll decrement k until it reaches 0 and at that point returns the value inside the current node.

The problem I have is that I'm not quite sure how to break out of the recursion of inorder traversal short of modifying the whole function, but I feel like having to modify the higher order function ruins the point of using a higher order function in the first place.

Is there a way to break after k iterations?

Answer 1

I observe that the results of the recursive call to go on the left and right subtrees are not available to f; hence no matter what f does, it cannot choose to ignore the results of recursive calls. Therefore I believe that inorder as written will always walk over the entire tree. (edit: On review, this statement may be a bit strong; it seems f may have a chance to ignore left subtrees. But the point basically stands; there is no reason to elevate left subtrees over right subtrees in this way.)

A better choice is to give the recursive calls to f. For example:

anyOldOrder :: (a -> b -> b -> b) -> b -> BinaryTree a -> b
anyOldOrder f z = go where
    go Leaf = z
    go (Node v l r) = f v (go l) (go r)

Now when we write

flatten = anyOldOrder (\v ls rs -> ls ++ [v] ++ rs) []

we will find that flatten is sufficiently lazy:

> take 3 (flatten (Node 'c' (Node 'b' (Node 'a' Leaf Leaf) Leaf) undefined))
"abc"

(The undefined is used to provide evidence that this part of the tree is never inspected during the traversal.) Hence we may write

findK k = take 1 . reverse . take k . flatten

which will correctly short-circuit. You can make flatten slightly more efficient with the standard difference list technique:

flatten' t = anyOldOrder (\v l r -> l . (v:) . r) id t []

Just for fun, I also want to show how to implement this function without using an accumulator list. Instead, we will produce a stateful computation which walks over the "interesting" part of the tree, stopping when it reaches the kth element. The stateful computation looks like this:

import Control.Applicative
import Control.Monad.State
import Control.Monad.Trans.Maybe

kthElem k v l r = l <|> do
    i <- get
    if i == k
        then return v
        else put (i+1) >> r

Looks pretty simple, hey? Now our findK function will farm out to kthElem, then do some newtype unwrapping:

findK' k = (`evalState` 1) . runMaybeT . anyOldOrder (kthElem 3) empty

We can verify that it is still as lazy as desired:

> findK' 3 $ Node 'c' (Node 'b' (Node 'a' Leaf Leaf) Leaf) undefined
Just 'c'

Answer 2

There are (at least?) two important generalizations of the notion of folding a list. The first, more powerful, notion is that of a catamorphism. The anyOldOrder of Daniel Wagner's answer follows this pattern.

But for your particular problem, the catamorphism notion is a bit more power than you need. The second, weaker, notion is that of a Foldable container. Foldable expresses the idea of a container whose elements can all be mashed together using the operation of an arbitrary Monoid. Here's a cute trick:

{-# LANGUAGE DeriveFoldable #-}

-- Note that for this trick  only I've 
-- switched the order of the Node fields.
data BinaryTree a =
  Node (BinaryTree a) a (BinaryTree a)
  | Leaf
  deriving (Show, Foldable)

index :: [a] -> Int -> Maybe a
[] `index` _ = Nothing
(x : _) `index` 0 = Just x
(_ : xs) `index` i = xs `index` (i - 1)

(!?) :: Foldable f => Int -> f a -> Maybe a
xs !? i = toList xs `index` i

Then you can just use !? to index into your tree!

---

That trick is cute, and in fact deriving Foldable is a tremendous convenience, but it won't help you understand anything. I'll start by showing how you can define treeToList fairly directly and efficiently, without using Foldable.

treeToList :: BinaryTree a -> [a]
treeToList t = treeToListThen t []

The magic is in the treeToListThen function. treeToListThen t more converts t to a list and appends the list more to the end of the result. This slight generalization turns out to be all that's required to make conversion to a list efficient.

treeToListThen :: BinaryTree a -> [a] -> [a]
treeToListThen Leaf more = more
treeToListThen (Node v l r) more =
  treeToListThen l $ v : treeToListThen r more

Instead of producing an inorder traversal of the left subtree and then appending everything else, we tell the left traversal what to stick on the end when it's done! This avoids the potentially serious inefficiency of repeated list concatenation that can turn things O(n^2) in bad cases.

Getting back to the Foldable notion, turning things into lists is a special case of foldr:

toList = foldr (:) []

So how can we implement foldr for trees? It ends up being somewhat similar to what we did with toList:

foldrTree :: (a -> b -> b) -> b -> BinaryTree a -> b
foldrTree _ n Leaf = n
foldrTree c n (Node v l r) = foldrTree c rest l
  where
    rest = v `c` foldrTree c n r

That is, when we go down the left side, we tell it that when it's done, it should deal with the current node and its right child.

Now foldr isn't quite the most fundamental operation of Foldable; that is actually

foldMap :: (Foldable f, Monoid m)
        => (a -> m) -> f a -> m

It is possible to implement foldr using foldMap, in a somewhat tricky fashion using a peculiar Monoid. I don't want to overload you with details of that right now, unless you ask (but you should look at the default definition of foldr in Data.Foldable). Instead, I'll show how foldMap can be defined using Daniel Wagner's anyOldOrder:

instance Foldable BinaryTree where
  foldMap f = anyOldOrder bin mempty where
    bin lres v rres = lres <> f v <> rres