CODEWITHSUNDEEP
haskelltreeindexing
rlhh
rlhh

Reputation: 903

Retrieving an element of a balanced binary tree in Haskell

Assuming I have a custom tree datatype of the following form:

data BalTree a = Leaf | Node Integer (BalTree a) a (BalTree a) deriving (Eq, Show, Read)

and creating a new tree of size 10, I'll get this:

Node 10 (Node 5 (Node 2 (Node 1 Leaf 'Z' Leaf) 'Z' Leaf)
                 'Z'
                (Node 2 (Node 1 Leaf 'Z' Leaf) 'Z' Leaf))
'Z'
        (Node 4 (Node 2 (Node 1 Leaf 'Z' Leaf) 'Z' Leaf)
                 'Z'
                 (Node 1 Leaf 'Z' Leaf))

How do I retrieve an element in in-order transversal when given an index?

My attempt: 

ind Leaf pos            = Nothing
ind tree@(Node n lt x rt) pos
    | pos < 0           = Nothing
    | pos > treeSize-1  = Nothing
    | pos < hTreeSize   = ind lt pos
    | pos == hTreeSize  = Just x
    | pos > hTreeSize   = ind rt (pos - hTreeSize)
    where treeSize = size tree
          hTreeSize = treeSize `div` 2

I'm not exactly sure if this is in-order transversal and it doesn't return the correct result.

Upvotes: 0

Views: 423

Answers (1)

scvalex
scvalex

Reputation: 15345

We want to get the nth value stored in a binary tree in an in-order walk. We know the number of values stored in each tree rooted at each node (the Integer parameter of Node).

data BalTree a = Leaf
               | Node Integer (BalTree a) a (BalTree a)

size :: BalTree a -> Integer
size Leaf              = 0
size (Node size _ _ _) = size

nthInOrder :: BalTree a -> Integer -> Maybe a
nthInOrder Leaf _ =
    Nothing
nthInOrder (Node _ left x right) n
    | leftSize == n - 1 = Just x
    | n <= leftSize     = nthInOrder left n
    | otherwise         = nthInOrder right (n - leftSize - 1)
  where
    leftSize  = size left

The idea is this: suppose we're at node A and want the nth value:

  A
 / \
B   C

If B holds n-1 values, then the nth value is that of A. If B holds more or equal than n values, then we can ignore the rest of the tree and search only B; so we just recurse into it. Otherwise, we should be looking for the value in C, so we recurse into it; in this case, we also need to update the n to reflect that there are some values in B, and 1 value in A.

In the worst case, this algorithm walks down to a Leaf, so, the complexity is O(depth of tree). If the tree is balanced, the complexity is O(log2(size of tree)).

Upvotes: 2

Related Questions