Data structure request: Lazily infinite set

Question

Are there F :: * -> *, iterate' :: Ord a => (a -> a) -> a -> F a and elem' :: Ord a => Int -> a -> F a -> Bool with the following properties?

• elem x (take n (iterate f y)) ⇒ elem' n x (iterate' f y) ⇒ elem x (iterate f y)

• elem' n x (iterate' f y) runs in O(n * log n) time and O(n) space

• elem' n x xs runs in O(log n) time and O(1) space

Answer 1

import qualified Data.Set as S

type F x = [S.Set x]

iterate' f
  = map head
  . evalState (traverse (state . splitAt) (iterate (*2) 1))
  . scanl (flip S.insert) S.empty
  . iterate f

elem' n x xs = S.member x $ xs !! (ceiling (logBase 2 (fromIntegral n)) - 1)

(Do the intermediate sets count as allocated space? Can you even do finite sets in linear space if you need to balance them?)