Skip to content

Instantly share code, notes, and snippets.

@Average-user
Last active April 21, 2020 16:22
Show Gist options
  • Select an option

  • Save Average-user/a875920b0ae6e6229808159392ab29f7 to your computer and use it in GitHub Desktop.

Select an option

Save Average-user/a875920b0ae6e6229808159392ab29f7 to your computer and use it in GitHub Desktop.
Implementation of the WROM algorithm for Constant Time generation of Free Trees
{-# LANGUAGE BangPatterns #-}
import qualified Data.Vector as V
import System.CPUTime
import System.Environment
import Text.Printf
{- Implementation of the WROM algorithm for finding all
free trees of a given order. The algorithm is explained
here:
https://www.cs.auckland.ac.nz/courses/compsci720s1c/lectures/mjd/treenotes.pdf
-}
type Tree = V.Vector Int
nextRootedTree :: Int -> Maybe Int -> Tree -> Tree
nextRootedTree n p' xs = s <$> V.generate n id
where
back = map ((,) <*> (V.!) xs) [n-1,n-2..0]
(p,lp) = case p' of
Nothing -> head $ filter ((/=1) . snd) back
(Just x) -> (x, xs V.! x)
q = fst $ head $ filter (\(i,e) -> i < p && e == lp-1) back
s i = if i < p then xs V.! i else s (i - p + q)
splitTree :: Int -> Tree -> (Tree, Tree)
splitTree n l = (l1, l2)
where
ops = V.filter ((==1) . snd) (V.indexed l)
m = if V.length ops < 2 then n else fst (ops V.! 1)
l1 = pred <$> (V.tail (V.take m l))
l2 = V.cons 0 (V.drop m l)
newCondition :: Ord a => a -> a -> Int -> Int -> Int -> Int -> Bool
newCondition l1 l2 hl1 hl2 ll1 ll2
| not (hl1 <= hl2) = False
| hl1 /= hl2 = hl1 <= hl2
| ll1 > ll2 = False
| ll1 == ll2 && l1 > l2 = False
| otherwise = hl1 <= hl2
nextTree :: Int -> Tree -> Tree
nextTree n l | newCondition l1 l2 hl1 hl2 ll1 ll2 = l
| (l V.! ll1) > 2 = newl'
| otherwise = newl
where
(l1,l2) = splitTree n l
(hl1,hl2) = (V.maximum l1, V.maximum l2)
(ll1,ll2) = (V.length l1, V.length l2)
newl = nextRootedTree n (Just ll1) l
(nl1,nl2) = splitTree n newl
nhl1 = V.maximum nl1
suffix = V.generate (nhl1 + 1) succ
newl' = V.take (n - nhl1 - 1) newl V.++ suffix
generateFree :: Int -> [Tree]
generateFree n = f $ V.fromList ([0..n `div` 2] ++ [1..(n+1) `div` 2 -1])
where
f xs | cond xs = [xs]
| otherwise = let nxs = nextTree n xs
in nxs : f (nextRootedTree n Nothing nxs)
cond v = V.head v == 0 && all (==1) (V.tail v)
-- Timing
main = do
n <- read . head <$> getArgs
putStrLn " Order | No. Of Trees | CPU Time | CPU Time per Tree"
putStrLn "-------|--------------|----------|------------------"
mapM_ (time (length . generateFree)) [1..n]
time :: (Int -> Int) -> Int -> IO ()
time f n = do
start <- getCPUTime
let !r = f n
end <- getCPUTime
let diff = (fromIntegral (end - start)) / (10^12) :: Double
perT = (diff / (fromIntegral r)) * (10^6)
putStrLn (printf "%6d | %12d | %6.1f s | %14.3f μs" n r diff perT)
{-
$ ./FreeTrees 20
Order | No. Of Trees | CPU Time | CPU Time per Tree
-------|--------------|----------|------------------
1 | 1 | 0.0 s | 0.000 μs
2 | 1 | 0.0 s | 0.000 μs
3 | 1 | 0.0 s | 0.000 μs
4 | 2 | 0.0 s | 0.000 μs
5 | 3 | 0.0 s | 0.000 μs
6 | 6 | 0.0 s | 0.000 μs
7 | 11 | 0.0 s | 0.000 μs
8 | 23 | 0.0 s | 0.000 μs
9 | 47 | 0.0 s | 0.000 μs
10 | 106 | 0.0 s | 0.000 μs
11 | 235 | 0.0 s | 0.000 μs
12 | 551 | 0.0 s | 7.260 μs
13 | 1301 | 0.0 s | 3.075 μs
14 | 3159 | 0.0 s | 1.266 μs
15 | 7741 | 0.0 s | 2.584 μs
16 | 19320 | 0.0 s | 2.277 μs
17 | 48629 | 0.1 s | 2.303 μs
18 | 123867 | 0.3 s | 2.260 μs
19 | 317955 | 0.8 s | 2.365 μs
20 | 823065 | 1.9 s | 2.367 μs
-}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment