module Tree
where

data Tree = Empty Int
          | Pair Int Tree Tree

index :: Tree -> Int
index (Empty n) = n
index (Pair n _ _) = n

data Result = HasLeft
            | HasRight
            | HasBoth Int
            | HasNeither



join :: Int -> Result -> Result -> Result
join _ (HasBoth k) _ = HasBoth k
join _ _ (HasBoth k) = HasBoth k
join _ HasNeither r = r
join _ r HasNeither = r
join k HasLeft HasRight = HasBoth k
join k HasRight HasLeft = HasBoth k


lca :: [Tree] -> Int -> Int -> Maybe Int
lca roots m n = answer (map ancestor roots)
  where extend :: Tree -> Result -> Result
        extend t = if index t == m then addLeft
                   else if index t == n then addRight
                   else id
           where addLeft HasRight    = HasBoth (index t)
                 addLeft HasNeither  = HasLeft
                 addRight HasLeft    = HasBoth (index t)
                 addRight HasNeither = HasRight
        ancestor t@(Empty _) = extend t HasNeither
        ancestor t@(Pair n left right) =
            extend t (join n (ancestor left) (ancestor right))
        answer :: [Result] -> Maybe Int
        answer (HasBoth n : _) = Just n
        answer (HasLeft : _) = Nothing
        answer (HasRight : _) = Nothing
        answer (HasNeither : results) = answer results
        answer [] = Nothing

tree0, tree1 :: Tree

{-

    0       1
   / \     / \
  2   3   4   5
 / \         / \
6   7       8   9
   / \
 10   11

-}

tree0 = Pair 0 two three
  where two = Pair 2 six seven
        three = Empty 3
        six = Empty 6
        seven = Pair 7 ten eleven
        ten = Empty 10
        eleven = Empty 11

tree1 = Pair 1 four five
    where four = Empty 4
          five = Pair 5 eight nine
          eight = Empty 8
          nine = Empty 9

forest :: [Tree]
forest = [tree0, tree1]


test = lca forest