Tree.hs

module Data.Tree where

import Data.List (sortBy)
import Data.Ord (comparing)

data Tree a = Tip | Node a (Tree a) (Tree a)
            deriving (Show, Eq)

-- x, y
type Coords = (Int, Int)

ex1 :: Tree Int
ex1 = Node 6
        (Node 3
          (Node 5
            (Node 9 Tip Tip)
            (Node 2
              Tip
              (Node 7 Tip Tip)))
          (Node 1 Tip Tip))
        (Node 4
          Tip
          (Node 0
             (Node 8 Tip Tip)
             Tip))

inorder :: Tree a -> [a]
inorder Tip = []
inorder (Node v l r) =
  inorder l ++ v : inorder r

columns :: Tree Int -> [Int]
columns = fromCoords . nodeCoords (0, 0)

nodeCoords :: Coords -> Tree Int -> Tree (Int, Coords)
nodeCoords _ Tip = Tip
nodeCoords (x, y) (Node v l r) = Node (v, (x, y)) (nodeCoords (x - 1, y + 1) l) (nodeCoords (x + 1, y + 1) r)

fromCoords :: Tree (Int, Coords) -> [Int]
fromCoords = map fst . sortBy (comparing snd) . inorder

TreeSpec.hs

{-# LANGUAGE ScopedTypeVariables #-}
module Data.TreeSpec (main, spec) where

import Test.Hspec
import Test.QuickCheck

import Data.Tree

-- `main` is here so that this module can be run from GHCi on its own.  It is
-- not needed for automatic spec discovery.
main :: IO ()
main = hspec spec

spec :: Spec
spec = do
  describe "inorder" $ do
    it "should return the inorder of a tree" $ do
      inorder ex1 `shouldBe` [ 9, 5, 2, 7, 3, 1, 6, 4, 8, 0 ]
  describe "columns" $ do
    it "should return the empty list for an empty tree" $
      columns Tip `shouldBe` []
    it "should return the single value of a one-node tree" $ property $
      \x -> columns (Node x Tip Tip) == [x]
    it "should return the values of a tree with two children" $ property $
      \x y z -> columns (Node x (Node y Tip Tip) (Node z Tip Tip)) == [y, x, z]
    it "should return the values of a tree with 'diamond' shape" $ property $
      \a b c d -> columns (Node a (Node b Tip (Node d Tip Tip))
                                  (Node c Tip Tip)) == [b, a, d, c]
    it "should return the values of a tree with three levels" $ property $
      \a b c d e f g -> let tree = Node a
                                     (Node b
                                           (Node d Tip Tip)
                                           (Node e Tip Tip))
                                     (Node c
                                           (Node f Tip Tip)
                                           (Node g Tip Tip))
                        in columns tree == [d, b, a, e, f, c, g]
    it "should return the desired result for the example tree" $
      columns ex1 `shouldBe`[ 9, 5, 3, 2, 6, 1, 7, 4, 8, 0 ]
  describe "nodeCoords" $ do
    it "should turn the empty tree into an empty tree" $ property $
      \x -> nodeCoords x Tip == (Tip :: Tree (Int, Coords))
    it "should turn a single leaf into a leaf at offset" $ property $
      \x a -> nodeCoords x (Node a Tip Tip) == Node (a, x) Tip Tip
    it "should turn a node with two children into -1 and 1 offset" $ property $
      \((x,y) :: (Int, Int)) a b c -> let tree = Node a (Node b Tip Tip) (Node c Tip Tip)
                                          tree' = Node (a, (x,y)) (Node (b, (x-1, y+1)) Tip Tip) (Node (c, (x+1, y+1)) Tip Tip)
                                      in nodeCoords (x,y) tree == tree'
  describe "fromCoords" $ do
    it "should return an empty list for an empty tree" $
      fromCoords Tip `shouldBe` []
    it "should return a single-element list for a 1-node tree" $ property $
      \x a -> fromCoords (Node (a, x) Tip Tip) == [a]