Solution to https://open.kattis.com/problems/ceiling, using paths to represent shapes of trees

CeilingFunctionAlt.hs

{-# LANGUAGE DeriveFoldable, DeriveFunctor, ScopedTypeVariables #-}

module Main where

import           Prelude                 hiding ( Left
                                                , Right
                                                )
import           Data.List               hiding ( insert )
import           Control.Monad

data BTree a = Leaf | Branch (BTree a) a (BTree a)
  deriving (Eq, Foldable, Functor, Show)

data Direction = Left | Right deriving (Eq, Ord, Show)

insert :: (Ord a) => a -> BTree a -> BTree a
insert x Leaf = Branch Leaf x Leaf
insert x (Branch leftys y rightys)
  | x < y     = Branch (insert x leftys) y rightys
  | otherwise = Branch leftys y (insert x rightys)

mkTree :: (Ord a) => [a] -> BTree a
mkTree = foldl (flip insert) Leaf

-- We can represent a binary tree exclusively by the set of
-- paths to its most removed (farthest) nodes.
-- For the sake of our representation, a tree which is a leaf
-- is represented by the empty list, while a tree without children
-- is represented by a list containing the empty list.
rep :: BTree a -> [[Direction]]
rep = go []
 where
  go ds (Branch Leaf   _ Leaf   ) = [ds]
  go ds (Branch leftys _ Leaf   ) = go (Left : ds) leftys
  go ds (Branch Leaf   _ rightys) = go (Right : ds) rightys
  go ds (Branch leftys _ rightys) =
    go (Left : ds) leftys ++ go (Right : ds) rightys
  go ds Leaf = []

countShapes :: [BTree a] -> Int
countShapes = length . group . sort . map rep

main :: IO ()
main = do
  -- The first line of the input contains two integers, n (1 <= n <= 50),
  -- which is the number of ceiling prototypes to analyze, and k
  -- (1 <= k <= 20), which is the number of layers in each of the prototypes.
  [n, k] :: [Int] <- map read . words <$> getLine
  -- The next n lines describe the ceiling prototypes. Each of these lines
  -- contains k distinct integers (between 1 and 10^6, inclusive), which
  -- are the collapse-resistance values of the layers in a ceiling prototype,
  -- ordered from top to bottom.
  rows :: [[Int]] <- map (map read . words) <$> replicateM n getLine
  let malformed = filter ((/= k) . length) rows
  when
    ((not . null) malformed)
    (  error
    $  "Expected rows of length "
    ++ show k
    ++ " but saw:\n"
    ++ show malformed
    )
  let trees = map mkTree rows
  (print . countShapes) trees