Nth element of a sorted tree

NthElement.hs

{-# LANGUAGE GADTs #-}

module NthElement where

import Data.Functor
import Control.Applicative
import Control.Monad
import Control.Monad.Trans.State

data Tree a where
  Nil  :: Tree a
  Node :: Ord a => Tree a -> a -> Tree a -> Tree a

getNthElement :: Int -> Tree a -> Either Int a
getNthElement k t = maybe (Left $ k - s) Right mv where

  (mv , s) = go t `runState` k

  go :: Tree a -> State Int (Maybe a)
  go Nil          = return Nothing
  go (Node l v r) = do
    rv <- go r
    k  <- get
    () <- put $ k - 1
    lv <- go l
    return $ rv <|> (v <$ guard (k == 0)) <|> lv

leftInfTree :: Tree Int
leftInfTree = foldr (\ k ih -> Node ih k Nil) Nil [1..]