sublinear indexing with trees

whatevs.hs

import Data.List
import System.Environment

data SizedTree a = Node Integer a [SizedTree a] deriving (Eq, Ord, Read, Show)
type Stream a = [SizedTree a]

indexT :: Integer -> SizedTree a -> a
indexT 0 (Node _ v _) = v
indexT n (Node _ _ vs) = indexS (n-1) vs

indexS :: Integer -> Stream a -> a
indexS n (t@(Node m _ _):ts)
	| n < m = indexT n t
	| otherwise = indexS (n-m) ts

fromListS :: [a] -> Stream a
fromListS = goS [1] where
	goS ns@(n:_) xs = goT ns b : goS (2*n+1:ns) e where
		(b, e) = genericSplitAt n xs

	goT [1] [x] = Node 1 x []
	goT (n:ns@(n':_)) (x:xs) = Node n x [goT ns b, goT ns e] where
		(b, e) = genericSplitAt n' xs

example :: Integer -> Integer
example n = sum [indexS i vs | i <- [0..n]] where vs = fromListS [0..]

main = getArgs >>= print . example . read . head

-- % /usr/bin/time ./test 100000
-- 5000050000
-- 0.19user 0.01system 0:00.21elapsed 99%CPU (0avgtext+0avgdata 25036maxresident)k
-- 0inputs+0outputs (0major+5507minor)pagefaults 0swaps
-- % /usr/bin/time ./test 1000000
-- 500000500000
-- 2.42user 0.06system 0:02.49elapsed 99%CPU (0avgtext+0avgdata 217408maxresident)k
-- 0inputs+0outputs (0major+53621minor)pagefaults 0swaps
-- /usr/bin/time ./test 1000000  2.43s user 0.06s system 99% cpu 2.499 total
-- % /usr/bin/time ./test 10000000
-- 50000005000000
-- 28.71user 0.63system 0:29.58elapsed 99%CPU (0avgtext+0avgdata 1906148maxresident)k
-- 0inputs+0outputs (0major+475784minor)pagefaults 0swaps
-- /usr/bin/time ./test 10000000  28.71s user 0.64s system 99% cpu 29.583 total