index function for balanced binary tree

I have problem, i can't figure out how i must decide what sub-tree my function indexJ must to choose at the each step walks through the my balanced binary tree - JoinList.

The idea is to cache the size (number of data elements) of each sub-tree. This can then be used at each step to determine if the desired index is in the left or the right branch.

I have this code:

data JoinList m a = Empty
                  | Single m a
                  | Append m (JoinList m a) (JoinList m a)
                  deriving (Eq, Show)

newtype Size = Size Int
  deriving (Eq, Ord, Show, Num)

getSize :: Size -> Int
getSize (Size i) = i

class Sized a where
  size :: a -> Size

instance Sized Size where
  size = id

instance Monoid Size where
  mempty  = Size 0
  mappend = (+)

i write functions:

tag :: Monoid m => JoinList m a -> m
tag Empty = mempty
tag (Single x dt) = x
tag (Append x l_list r_list) = x

(+++) :: Monoid m => JoinList m a -> JoinList m a -> JoinList m a
(+++) jl1 jl2 = Append (mappend (tag jl1) (tag jl2)) jl1 jl2

indexJ :: (Sized b, Monoid b) => Int -> JoinList b a -> Maybe a
indexJ _ Empty = Nothing
indexJ i jl | i < 0 || (i+1) > (sizeJl jl) = Nothing 

  where sizeJl = getSize . size . tag

indexJ 0 (Single m d) = Just d
indexJ 0 (Append m (Single sz1 dt1) jl2) = Just dt1
indexJ i (Append m jl1 jl2) = if (sizeJl jl1) >= (sizeJl jl2) 
                              then indexJ (i-1) jl1  
                              else indexJ (i-1) jl2 

  where sizeJl = getSize . size . tag

functions tag and (+++) working well, but i need to finish indexJ function, which must return i-th element from my JoinList tree, i = [0..n]

my function indexJ working wrong =) if i have empty tree - it's (Size 0) if i have Single (Size 1) "data" - it's (Size 1) but what about if i have Append (Size 2) (Single (Size 1) 'k') (Single (Size 1) 'l') what branch i must choose? i-1 = 1 and i have two branches with 1 data element in each.

UPDATE: if someone needs take and drop functions for JoinList's trees i make it:

dropJ :: (Sized b, Monoid b) => Int -> JoinList b a -> JoinList b a
dropJ _ Empty = Empty 
dropJ n jl | n <= 0 = jl
dropJ n jl | n >= (getSize . size $ tag jl) = Empty
dropJ n (Append m jL1 jL2)
  | n == s1 = jL2
  | n < s1 = (dropJ n jL1) +++ jL2
  | otherwise = dropJ (n - s1) jL2
                where s1 = getSize . size $ tag jL1

takeJ :: (Sized b, Monoid b) => Int -> JoinList b a -> JoinList b a
takeJ _ Empty = Empty 
takeJ n jl | n <= 0 = Empty
takeJ n jl | n >= (getSize . size $ tag jl) = jl
takeJ n (Append m jL1 jL2)
  | n == s1 = jL1
  | n < s1 = (takeJ n jL1)
  | otherwise = jL1 +++ takeJ (n - s1) jL2
                where s1 = getSize . size $ tag jL1

Upvotes: 2

Views: 410

Answers (2)

Daniel Fischer
Daniel Fischer

Reputation: 183878

I suppose in

Append m joinList1 joinList2

the elements of joinList1 are meant to take up the first indices, followed by the elements of joinList2.

Then, when indexing

indexJ i (Append m jL1 jL2)

you have to compare i to the size of jL1 - let us call that s1. Then the elements of jL1 occupy the indices 0 to s1 - 1, and the elements of jL2 occupy the indices from s1 to s1 + s2 - 1, hence

indexJ :: (Sized b, Monoid b) => Int -> JoinList b a -> Maybe a
indexJ _ Empty  = Nothing
indexJ i (Single m d)
    | i == 0    = Just d
    | otherwise = Nothing
indexJ i (Append m jL1 jL2)
    | i < 0     = Nothing
    | i >= getSize (size m) = Nothing     -- optional, more efficient to have it
    | i < s1    = indexJ i jL1
    | otherwise = indexJ (i - s1) jL2
      where
        s1 = getSize . size $ tag jL1

if the index is smaller than s1, we look in the first sublist, otherwise in the second.

Upvotes: 6

chris
chris

Reputation: 5018

Typically you would encode the position in a tree-structure by sequences of numbers, not just a single number. For example (assuming indices start at 0):

[] -- empty sequence = root of tree
[0,1] -- first follow the first child, then the second child
[0,0,0] -- go 3 levels down in the leftmost branch

That would make the implementation of an index function much simpler.

Upvotes: 1

Related Questions