FullBinaryTree.hs

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NPlusKPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DeriveFunctor #-}

module FullBinaryTree where

import Control.Arrow ((&&&))
import Control.Comonad.Cofree (Cofree (..))
import Data.Functor.Base (NonEmptyF (..))
import Data.Functor.Foldable
import Data.List.NonEmpty hiding (zipWith, reverse)
import Numeric.Natural

-- import Catalan

data TreeF a
  = LeafF
  | a :^^: a
  deriving Functor

data Tree
  = Leaf
  | Tree :^: Tree

type instance Base Tree = TreeF

instance Recursive Tree where
  project = \ case
    Leaf -> LeafF
    s :^: t -> s :^^: t

instance Corecursive Tree where
  embed = \ case
    LeafF    -> Leaf
    s :^^: t -> s :^: t

(×) :: [Tree] -> [Tree] -> [Tree]
ts × us = [ t :^: u | t <- ts, u <- us]

nTrees :: Natural -> [Tree]
nTrees = dyna trees natural
  where
    trees :: NonEmptyF Natural (Cofree (NonEmptyF Natural) [Tree]) -> [Tree]
    trees = \ case
      NonEmptyF 0 Nothing      -> [Leaf]
      NonEmptyF n (Just table) -> concat (zipWith (×) xs (reverse xs))
        where
          xs = tk n table

nTreesSpec :: Natural -> [Tree]
nTreesSpec = \ case
  0   -> [Leaf]
  n+1 -> concat $ zipWith (×) <*> reverse $ nTreesSpec <$> [0 .. n]

size :: Tree -> Natural
size = cata phi
  where
    phi = \ case
      LeafF    -> 0
      m :^^: n -> succ (m + n)

catalan :: Natural -> Natural
catalan = iter 1 0
  where
    iter c _ 0 = c
    iter c m n = iter (nextCatalan c m) (succ m) (pred n)
      where

nextCatalan :: Natural -> Natural -> Natural
nextCatalan c n = c * (4 * n + 2) `div` (n + 2)

catalans :: [Natural]
catalans@(_:catalans') = 1 : zipWith nextCatalan catalans [0 ..]

catalanAcc :: Natural -> Natural
catalanAcc = iter 0 1 0
  where
    iter a _ _ 0 = a
    iter a c m n = iter (a + c) (nextCatalan c m) (succ m) (pred n)

---

dyna :: Functor f => (f (Cofree f a) -> a) -> (c -> f c) -> (c -> a)
dyna phi psi = hd . h
  where
    h = uncurry (:<) . (phi &&& id) . (h <$>) . psi

hd :: Cofree f a -> a
hd = \ case
  x :< _ -> x

tk :: Natural -> (Cofree (NonEmptyF b) a) -> [a]
tk 0     _ = []
tk (n+1) (a :< NonEmptyF _ Nothing)   = [a]
tk (n+1) (a :< NonEmptyF _ (Just as)) = a : tk n as

natural :: Natural -> NonEmptyF Natural Natural
natural = \ case
  0   -> NonEmptyF 0 Nothing
  n+1 -> NonEmptyF (n+1) (Just n)