Search code examples
haskelltime-complexitybinary-treerecursion-schemes

Selectively recurse into left or right subtree of a binary tree using a catamorphism (or any recursion scheme)


I'm trying to implement a binary search tree (or set) using fixed points of functors. I've defined my fixed point as follows:

newtype Fix f = In (f (Fix f))

out :: Fix f -> f (Fix f)    
out (In f) = f    
    
-- Catamorphism    
type Algebra f a = f a -> a    
    
cata :: (Functor f) => Algebra f a -> Fix f -> a    
cata f = f . fmap (cata f) . out 

To make the binary tree, I'm using a red-black tree like so:

data NodeColor = Red | Black deriving (Eq, Show)    
    
data RedBlackTreeF a r = EmptyRedBlackTreeF | RedBlackTreeNodeF NodeColor r a r deriving (Eq, Show)
    
instance Functor (RedBlackTreeF a) where                                             
        fmap _ EmptyRedBlackTreeF = EmptyRedBlackTreeF                               
        fmap f (RedBlackTreeNodeF color r1 a r2) =                                   
                RedBlackTreeNodeF color (f r1) a (f r2)                              
                                                                                     
type RedBlackTreeF' a = Fix (RedBlackTreeF a) 

The traditional benefit of a binary tree is being able to cut down search time by choosing whether to search further in the left or right subtree like so (in psuedocode):

fun member (x, E) = false
   | member (x, T (_, a, y, b)) =
     if x < y then member (x, a)
     else if x > y then member (x, b)
     else true

The function member will go left if the element that is being searched for is less than the current element and right if the opposite is true. It therefore improves search time to O(logn).

However, in a recursion scheme, the algebra is recursively applied to the entire data structure. I've written an member algebra here:

memberPAlg :: Ord a => a -> RedBlackTreeF a Bool -> Bool    
memberPAlg _ EmptyRedBlackTreeF = False    
memberPAlg elem (RedBlackTreeNodeF _ left cur right) =    
        (elem == cur) || (left || right) 

But it seems to be O(nlogn) rather than O(logn). Is there any way to selectively recurse using a recursion scheme to save time complexity? Am I thinking about this the wrong way?


Solution

  • Because of laziness, left and right are evaluated only if you ask for them. So, just compare the current node with the value being searched for to decide which subtree to go into:

    memberPAlg :: Ord a => a -> RedBlackTreeF a Bool -> Bool    
    memberPAlg _ EmptyRedBlackTreeF = False    
    memberPAlg elem (RedBlackTreeNodeF _ left cur right) =    
       case compare elem cur of
         EQ -> True
         LT -> left
         GT -> right