import Data.Map as Map
import Data.Maybe
import Control.Applicative

data Element = A | B | C | D
    deriving (Eq, Ord, Show)

addElement :: Element -> Map Element Int -> Map Element Int
addElement = alter ((<|> Just 1) . fmap (+1))

getScore :: Element -> Map Element Int -> Float
getScore e m = fromMaybe 1.0 $ ((1.0/) . fromIntegral) <$> Map.lookup e m

data BinaryTree a = Empty | Node a (BinaryTree a) (BinaryTree a) deriving Show
type ElementTree = BinaryTree Element
type ScoredElementTree = BinaryTree (Element, Float)

assignScores :: ElementTree -> Map Element Int -> (Map Element Int, ScoredElementTree)
assignScores Empty m = (Map.empty, Empty)
assignScores (Node e t1 t2) m = (m', Node (e, getScore e m) t1' t2')
    where
        (m1, t1') = assignScores t1 m
        (m2, t2') = assignScores t2 m
        m' = addElement e $ unionWith (+) m1 m2

doAssignScores :: ElementTree -> ScoredElementTree
doAssignScores t = t'
    where (m, t') = assignScores t m