!module LOT

!import "lib/base.tri" Lib

main = exampleTwo
-- Level Order Traversal of a labelled binary tree
-- Objective: Print each "level" of the tree on a separate line
--
-- We model labelled binary trees as nested lists where values act as labels. We
-- require explicit notation of empty nodes. Empty nodes can be represented
-- with an empty list, `[]`, which evaluates to a single node `t`.
--
-- Example tree inputs:
-- [("1") [("2") [("4") t t] t] [("3") [("5") t t] [("6") t t]]]]
-- Graph:
--       1
--      / \
--     2   3
--    /   /  \
--   4   5    6

label = \node : Lib.head node

left = (\node : Lib.if (Lib.emptyList? node)
  [] 
  (Lib.if (Lib.emptyList? (Lib.tail node)) 
    [] 
    (Lib.head (Lib.tail node))))

right = (\node : Lib.if (Lib.emptyList? node) 
  [] 
  (Lib.if (Lib.emptyList? (Lib.tail node)) 
    [] 
    (Lib.if (Lib.emptyList? (Lib.tail (Lib.tail node))) 
      [] 
      (Lib.head (Lib.tail (Lib.tail node))))))

processLevel = Lib.y (\self queue : Lib.if (Lib.emptyList? queue) 
  [] 
  (Lib.pair (Lib.map label queue) (self (Lib.filter 
    (\node : Lib.not? (Lib.emptyList? node)) 
      (Lib.lconcat (Lib.map left queue) (Lib.map right queue))))))

levelOrderTraversal_ = \a : processLevel (t a t)

toLineString = Lib.y (\self levels : Lib.if (Lib.emptyList? levels) 
  "" 
  (Lib.lconcat 
    (Lib.lconcat (Lib.map (\x : Lib.lconcat x " ") (Lib.head levels)) "") 
    (Lib.if (Lib.emptyList? (Lib.tail levels)) "" (Lib.lconcat (t (t 10 t) t) (self (Lib.tail levels))))))

levelOrderToString = \s : toLineString (levelOrderTraversal_ s)

flatten = Lib.foldl (\acc x : Lib.lconcat acc x) ""

levelOrderTraversal = \s : Lib.lconcat (t 10 t) (flatten (levelOrderToString s))

exampleOne = levelOrderTraversal [("1") 
                                 [("2") [("4") t t] t] 
                                 [("3") [("5") t t] [("6") t t]]]

exampleTwo = levelOrderTraversal [("1") 
                                 [("2") [("4") [("8") t t] [("9") t t]] 
                                        [("6") [("10") t t] [("12") t t]]] 
                                 [("3") [("5") [("11") t t] t] [("7") t t]]]