select :: [a] -> [(a,[a])]
select [] = []
select (x:xs) = (x,xs) : [(y,x:ys) | (y,ys) <- select xs]

row :: [Int] -> [[Int]] -> [[Int]]
row [] [] = [[]]
row remaining (used:otherColumns) =
  [x:rest | (x,available) <- select remaining,
            x `notElem` used,
            rest <- row available otherColumns]
row _ _ = []

latinPruned :: Int -> [[[Int]]]
latinPruned n
  | n < 0     = error "negative order"
  | otherwise = build n (replicate n [])
  where
    build 0 usedColumns = [[]]
    build remainingRows usedColumns =
      [r:rest | r <- row [1..n] usedColumns,
                rest <- build (remainingRows-1)
                              (zipWith (:) r usedColumns)]

squaresRows :: Int -> [[[Int]]]
squaresRows n
  | n < 0     = error "negative order"
  | otherwise = cp (replicate n (perms [1..n]))

latinRows :: Int -> [[[Int]]]
latinRows = filter valid . squaresRows

valid :: [[Int]] -> Bool
valid rows =
  let n = length rows
      symbols = [1..n]
      goodRows = all (\r -> sort r == symbols) rows
      goodColumns = all (\j -> sort [r !! j | r <- rows] == symbols)
                        [0..n-1]
  in goodRows && goodColumns

latin :: Int -> [[[Int]]]
latin = filter valid . squares

squares :: Int -> [[[Int]]]
squares n
  | n < 0     = error "negative order"
  | otherwise = cp (replicate n (cp (replicate n [1..n])))

cp :: [[a]] -> [[a]]
cp = foldr extend [[]]
  where
    extend choices suffixes = [x:ys | x <- choices, ys <- suffixes]

unfoldTree :: (s -> Maybe (s,[a],s)) -> s -> Tree a
unfoldTree split seed =
  case split seed of
    Nothing -> Empty
    Just (left,equal,right) ->
      Node (unfoldTree split left) equal (unfoldTree split right)

toTreeUnfold :: Ord a => [a] -> Tree a
toTreeUnfold = unfoldTree split
  where
    split [] = Nothing
    split (pivot:xs) =
      let (lo,eq,hi) = partition3 pivot xs
      in Just (lo,pivot:eq,hi)

foldTree :: (b -> [a] -> b -> b) -> b -> Tree a -> b
foldTree combine zero Empty = zero
foldTree combine zero (Node left equal right) =
  combine (foldTree combine zero left)
          equal
          (foldTree combine zero right)

fromTreeFold :: Tree a -> [a]
fromTreeFold = foldTree (\left equal right -> left ++ equal ++ right) []

data Tree a = Empty | Node (Tree a) [a] (Tree a)

toTree :: Ord a => [a] -> Tree a
toTree [] = Empty
toTree (pivot:xs) =
  let (lo,eq,hi) = partition3 pivot xs
  in Node (toTree lo) (pivot:eq) (toTree hi)

fromTree :: Tree a -> [a]
fromTree Empty = []
fromTree (Node left equal right) =
  fromTree left ++ equal ++ fromTree right

partition3 :: Ord a => a -> [a] -> ([a],[a],[a])
partition3 pivot = foldr step ([],[],[])
  where
    step x rest =
      let (lo,eq,hi) = rest
      in case compare x pivot of
           LT -> (x:lo,eq,hi)
           EQ -> (lo,x:eq,hi)
           GT -> (lo,eq,x:hi)

sort :: Ord a => [a] -> [a]
sort [] = []
sort (pivot:xs) =
  let (lo,eq,hi) = partition3 pivot xs
  in sort lo ++ (pivot:eq) ++ sort hi

spanFold :: (a -> Bool) -> [a] -> ([a], [a])
spanFold p = foldr step ([], [])
  where
    step x ~(ys,zs)
      | p x       = (x:ys,zs)
      | otherwise = ([],x:(ys ++ zs))

spanP :: (a -> Bool) -> [a] -> ([a], [a])
spanP p [] = ([], [])
spanP p whole@(x:xs)
  | p x       = let (ys,zs) = spanP p xs in (x:ys,zs)
  | otherwise = ([],whole)

takeP :: (a -> Bool) -> [a] -> [a]
takeP p = foldr step []
  where
    step x rest = if p x then x : rest else []