Solution for https://www.geeksforgeeks.org/averages-levels-binary-tree/

Main.hs

module Main where

import Data.Maybe
import Data.List

-- Solution for https://www.geeksforgeeks.org/averages-levels-binary-tree/

data Tree a
  = Node a (Tree a) (Tree a)
  | Leaf
  deriving (Show)

--     4
--    / \
--   2   9
--  / \   \
-- 3   5   7

testTree :: Tree Float
testTree =
  Node
    4
    (Node 2 (Node 3 Leaf Leaf) (Node 5 Leaf Leaf))
    (Node 9 Leaf (Node 7 Leaf Leaf))

value :: Tree a -> Maybe a
value Leaf = Nothing
value (Node x _ _) = Just x

children :: Tree a -> [Tree a]
children (Leaf) = []
children (Node _ l r) = [l, r]

nextLevel :: [Tree a] -> [Tree a]
nextLevel = concatMap children

average :: [Float] -> Float
average xs = sum xs / genericLength xs

levels :: [Tree a] -> [[a]]
levels = takeWhile (not . null) . map (mapMaybe value) . iterate nextLevel

main :: IO ()
main = print $ map average $ levels [testTree]