data Strom a
  = Vetev (Strom a)
          a
          (Strom a)
  | Konec
  deriving (Show)

test = Vetev (Vetev Konec 1 Konec) 2 Konec

instance Eq a => Eq (Strom a) where
  Konec == Konec = True
  Vetev a b c == Vetev d e f = b == e && a == d && c == f
  _ == _ = False

koren (Vetev _ b _) = Just b
koren Konec = Nothing

vetve (Konec) = []
vetve (Vetev a _ c) = [a, c]

strom2seznam a = takeWhile (not . null) $ iterate (concat . map vetve) [a]

seznam2koren a = map (map koren) a

--reseni z hodiny, vyrobi jedno extra prazdne posledni patro
patra = map catMaybes . seznam2koren . strom2seznam

--moje reseni (to na konci je "angry monkey operator")
patra' =
  takeWhile (not . null) .
  map (catMaybes . map koren) .
  iterate (concatMap vetve) .
  (:[])

--unixove reseni
infixl 9 & --deklaruje operator &, coz budeme pouzivat jako unixove |

(&) = flip (.)

patraForstStyle =
  (:[]) &
  iterate (concatMap vetve) &
  map (map koren & catMaybes) &
  takeWhile (null & not)

--oficialni catMaybes existuje v Data.Maybe
catMaybes :: [Maybe a] -> [a]
catMaybes a = [b | Just b <- a]
 
--Strom jde poskladat v nejakem jednoduchem poradi, takze je Foldable
instance Foldable Strom where
  foldr f a (Konec) = a
  foldr f a (Vetev l c r) =
    let a1 = foldr f a r
        a2 = f c a1
     in foldr f a2 l

--Strom je kontejnerovity, takze je Functor
instance Functor Strom where
  fmap f Konec = Konec
  fmap f (Vetev a b c) = Vetev (fmap f a) (f b) (fmap f c)

expandujStrom s =
  let ss = sum s
   in Vetev s ss $ fmap (+ ss) s