saw.hs

{-# LANGUAGE BangPatterns #-}

import Data.Maybe
import Data.Array
import Data.Int

data Relation = G | L | N deriving (Show, Enum, Ord, Eq, Bounded)
data SCond = Any | F !Relation !Int64 deriving (Show, Ord, Eq)

key (Any,   n) = (-1, -1, n)
key (F r p, n) = (fromEnum r, p, n)
minKey = (-1, -1, -1)
maxKey !k !n = (fromEnum (maxBound :: Relation), k, n)

toFun :: SCond -> (Int64 -> Bool)
toFun Any      = \_ -> True
toFun (F N !p) = (/=p)
toFun (F L !p) = (< p)
toFun (F G !p) = (> p)

genNextCond :: SCond -> Int64 -> SCond
genNextCond Any !h = (F N h)
genNextCond (F _ p1) p2 | p1 < p2 = (F L p2)
                        | p1 > p2 = (F G p2)

saw ::  Int64 -> Int64 -> Integer
saw n k = sawl Any n
  where 
    sawl p n = cache!(key (p, n)) 
    cache = array (minKey, maxKey k n) [ (key k, saw' p n) | k@(p,n) <- genKeys ]
    --
    saw'  _  0 = 1
    saw' !p !n = sum [ sawl (genNextCond p h) (n-1) | h <- heads p]
    --
    heads !p = filter (toFun p) alpha
    alpha    = [1..k]
    --
    genKeys = (Any, n) : [ (c, nn) | nn <- [0..n], c <- conds ]
      where
        conds = [ (F f p) | p <- alpha, f <- relations ]
        relations = enumFromTo minBound maxBound

main :: IO ()
main = do
  args   <- getLine
  [n, k] <- return $ map read $ words args
  putStrLn $ show $ saw n k