FPChallenge December 2014

bugs.hs

import UU.PPrint hiding ((<$>))
import Data.List (nub, sortBy, elemIndex, groupBy)
import Data.Function (on)
import Data.Functor ((<$>))
import Data.Maybe (fromJust, fromMaybe)
import Control.Arrow (first, second)
import GHC.Exts (the)

data Tree a b = Tree a [(b, Tree a b)]
  
decisionTree :: [String] -> [(String, [String])] -> Tree [String] String
decisionTree [name] _ = Tree [name] []
decisionTree names [] = Tree names []
decisionTree names rows 
  | power answers == 1 || null forest' = decisionTree names $ tail rows
  | otherwise = Tree [question] $ forest' ++ [defaultCase]
  where
    power = length.nub
    sortedRows = sortBy (flip compare `on` power.snd) rows
    (question, answers) = head sortedRows
    splitted = first the . unzip <$> (groupBy ((==) `on` fst) $ sortBy (compare `on` fst) $ zip answers names)
    splitted' = second (++ fromMaybe [] ("—" `lookup` splitted)) <$> filter ((/="—").fst) splitted
    takenIndexes = fmap $ fromJust.(`elemIndex` names)
    takeOnly is xs = foldr (\i acc -> (xs !! i) : acc) [] is
    newRows subnames = second (takeOnly $ takenIndexes subnames) <$> tail sortedRows
    forest = second (\newNames -> (newNames, decisionTree newNames (newRows newNames))) <$> splitted'
    forest' = filter ((/="—").fst) $ second snd <$> filter ((length names/=).length.fst.snd) forest
    defaultCase = ("—", decisionTree names $ tail sortedRows)

inputMatrix :: ([String], [(String, [String])])
inputMatrix = (["Аурата сетуньская", "Десятилиньята лепая", "Семипунктата Коха", "Популий грыжомельский", "Гортикола филоперьевая"],
              [("Наличие бомбурий",         ["Да",      "Да",        "Нет",       "Да",    "Нет"]),
               ("Количество клептиконов",   ["1",       "1",         "0",         "3",     "5"]),
               ("Цвет велория",             ["Красный", "Оранжевый", "Оранжевый", "—",     "Синий"]),
               ("Наличие пумпеля",          ["Нет",     "Да",        "Да",        "—",     "—"]),
               ("Величина пумпеля",         ["—",       "Большой",   "Маленький", "—",     "—"]),
               ("Возможность крокотания",   ["Нет",     "Нет",       "—",         "Да",    "Нет"]),
               ("Возможность бульботания",  ["Нет",     "Да",        "—",         "Да",    "Нет"]),
               ("Наличие дуков и труков",   ["—",       "—",         "—",         "—",     "Да"]),
               ("Цвет лемпелей",            ["Жёлтый",  "Жёлтый",    "Жёлтый",    "Белый", "Белый"]),
               ("Наличие пильских трапков", ["Да",      "Да",        "Да",        "Да",    "Да"])])
-- Ввод-вывод
-- ==================================================================================

show' = show.fmap pretty

main :: IO ()
main = do
  let tree = uncurry decisionTree inputMatrix
  writeFile "solution_tree.txt" $ drawTreeVertical tree
  putStrLn "Дерево решений сохранено в solution_tree.txt"
  interactive tree                                                                       

interactive :: Tree [String] String -> IO ()
interactive (Tree [name] []) = putStrLn $ "Жук определен: " ++ name
interactive (Tree names []) = putStrLn $ "Жук достоверно не определен, варианты: " ++ show' names
interactive (Tree [question] forest) = do 
  putStrLn question
  putStrLn $ "Варианты ответа: " ++ show' answers
  ans <- untilM (`elem` answers) getLine
  interactive $ branches !! fromJust (ans `elemIndex` answers)
  where
    answers = (\x->if x == "—" then "Не знаю" else x) <$> fst <$> forest
    branches = snd <$> forest

untilM :: (Monad m) => (a -> Bool) -> m a -> m a
untilM p m =
   let go =
         do x <- m
            if p x
              then return x
              else go
   in  go

mapWithFirstLast :: (Bool -> Bool -> a -> b) -> [a] -> [b]
mapWithFirstLast f = go True where
  go b (x : []) = f b True  x : []
  go b (x : xs) = f b False x : go False xs

drawTreeVertical :: Tree [String] String -> String
drawTreeVertical tree = unlines (go tree) where
  go :: Tree [String] String -> [String]
  go (Tree x cs) = case cs of
    [] -> ["-- " ++ show' x]
    _  -> concat $ mapWithFirstLast (f (show' x)) $ map (second go) cs
    
  f :: String -> Bool -> Bool -> (String, [String]) -> [String] 
  f label bf bl (brLabel, (l:ls)) =
        let 
            spaces = map (const ' ') 
            dashes = map (const '-')
            indent = if bl then "  " ++spaces label++"  " else  " |" ++ spaces (label ++ brLabel) ++ "  "
            gap    = if bl then []                        else [" |" ++ spaces (label ++ brLabel) ++ "  "]
            branch = if bl && not bf
                           then " \\"++dashes label++"--" ++ brLabel
                           else if bf 
                             then "-(" ++ label        ++ ")-" ++ brLabel
                             else " +" ++ dashes label ++ "--" ++ brLabel
        in  (branch++l) : map (indent++) ls ++ gap

Из внешних библиотек используется только uulib для печати UTF8 строк (стандартный show все символы экранировал)