Pearls of Functional Programming Chap. 7 code

pearls-ch7.hs

#!/usr/bin/env stack
-- stack --install-ghc runghc --package=containers --package=pretty-tree --package=criterion --package=QuickCheck

-- Pearls of Functional Algorithm Design, Chap. 7
-- Building a tree of minimum height

{-# LANGUAGE DeriveGeneric, DeriveAnyClass #-}

import Control.DeepSeq (NFData, deepseq)
import Criterion.Main
import Data.Function ((&))
import Data.List (intercalate)
import GHC.Generics (Generic)
import qualified Data.Tree
import qualified Data.Tree.Pretty
import Test.QuickCheck (Gen, generate, resize, vectorOf, arbitrary)

data Tree = Leaf Int | Fork Tree Tree
    deriving (Show, Eq, Generic, NFData)

showTree :: Tree -> String
showTree = removeExtraPipes . Data.Tree.Pretty.drawVerticalTree . convertTree
    where
        convertTree :: Tree -> Data.Tree.Tree String
        convertTree (Leaf a) = Data.Tree.Node (show a) []
        convertTree (Fork t u) =
            let t' = convertTree t
                u' = convertTree u
            in
                Data.Tree.Node "|" [t', u']
        removeExtraPipes :: String -> String
        removeExtraPipes =
            unlines
            . map snd
            . filter (\(n, s) -> n `mod` 4 /= 1)
            . zip [0..]
            . lines

printTree :: Tree -> IO ()
printTree = putStr . showTree

printTrees :: [Tree] -> IO ()
printTrees = putStr . intercalate "\n" . map showTree

cost :: Tree -> Int
cost (Leaf x) = x
cost (Fork u v) = 1 + (cost u `max` cost v)

minBy :: Ord a => (b -> a) -> [b] -> b
minBy f = foldl1 (cmp f)

cmp :: Ord a => (b -> a) -> b -> b -> b
cmp f u v = if f u <= f v then u else v

mincostTree1 :: [Int] -> Tree
mincostTree1 = minBy cost . trees1

trees1 :: [Int] -> [Tree]
trees1 [x] = [Leaf x]
trees1 (x : xs) = concatMap (prefixes1 x) (trees1 xs)

-- Shows the trees that could result from inserting a new Int
-- as a leftmost node in an existing tree.
prefixes1 :: Int -> Tree -> [Tree]
prefixes1 x t@(Leaf y) = [Fork (Leaf x) t]
prefixes1 x t@(Fork u v) =
    [Fork (Leaf x) t]
    ++ [Fork u' v | u' <- prefixes1 x u]

-- A type of fold for nonempty lists. Accepts a function to apply to a singleton,
-- and a function to combine the next element with the prior result.
foldrn :: (a -> b -> b) -> (a -> b) -> [a] -> b
foldrn f g [x] = g x
foldrn f g (x : xs) = f x (foldrn f g xs)

-- A shorter version of trees defined in terms of foldrn
trees2 :: [Int] -> [Tree]
trees2 = foldrn (concatMap . prefixes1) (wrap . Leaf)

wrap :: a -> [a]
wrap x = [x]

type Forest = [Tree]

trees3 :: [Int] -> Forest
trees3 = map rollup . forests

forests :: [Int] -> [Forest]
forests = foldrn (concatMap . prefixes3) (wrap . wrap . Leaf)

prefixes3 :: Int -> Forest -> [Forest]
prefixes3 x ts = [Leaf x : rollup (take k ts) : drop k ts | k <- [1..length ts]]

rollup :: Forest -> Tree
rollup = foldl1 Fork

mincostTree4 :: [Int] -> Tree
mincostTree4 = foldl1 Fork . map snd . foldrn insert (wrap . leaf)

insert :: Int -> [(Int, Tree)] -> [(Int, Tree)]
insert x ts = leaf x : split x ts

split :: Int -> [(Int, Tree)] -> [(Int, Tree)]
split x [u] = [u]
split x (u : v : ts) =
    if x `max` fst u < fst v
        then u : v : ts
        else split x (fork u v : ts)

leaf :: Int -> (Int, Tree)
leaf x = (x, Leaf x)

fork :: (Int, Tree) -> (Int, Tree) -> (Int, Tree)
fork (a, u) (b, v) = (1 + a `max` b, Fork u v)

-- Useful for debugging/visualizing
spine :: Tree -> Forest
spine = reverse . revSpine
    where
        revSpine (Leaf x) = [(Leaf x)]
        revSpine (Fork x y) = y : revSpine x

main = do
    i100 <- generate $ genInts 100
    i300 <- generate $ genInts 300
    i1000 <- generate $ genInts 1000
    i3000 <- generate $ genInts 3000
    deepseq [i100, i300, i1000] $ return ()
    defaultMain
        [ bgroup "input size"
            [ bench "100 ints" $ nf mincostTree4 i100
            , bench "300 ints" $ nf mincostTree4 i300
            , bench "1000 ints" $ nf mincostTree4 i1000
            , bench "3000 ints" $ nf mincostTree4 i3000
            ]
        ]

genInts :: Int -> Gen [Int]
genInts n = vectorOf n $ fmap (\x -> abs x + 1) $ resize (n*2) $ arbitrary