Cartesian Tree

CartesianTree.hs

{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ExistentialQuantification #-}

module CartesianTree where

import Control.Lens
import Control.Monad
import Data.List (sortBy)
import Data.Maybe
import Data.Ord
import System.Random

class CartesianTree c k v p where
  priority ::
       forall f. Applicative f
    => (p -> f p)
    -> c k v p
    -> f (c k v p)
  key ::
       forall f. Applicative f
    => (k -> f k)
    -> c k v p
    -> f (c k v p)
  value ::
       forall f. Applicative f
    => (v -> f v)
    -> c k v p
    -> f (c k v p)
  left ::
       forall f. Applicative f
    => (c k v p -> f (c k v p))
    -> c k v p
    -> f (c k v p)
  right ::
       forall f. Applicative f
    => (c k v p -> f (c k v p))
    -> c k v p
    -> f (c k v p)
  nil ::
       forall f. Applicative f
    => f (c k v p)
  createNode ::
       forall f. Applicative f
    => k
    -> v
    -> p
    -> c k v p
    -> c k v p
    -> f (c k v p)

data CTree k v p
  = CNil
  | CNode { _ckey :: !k
          , _cvalue :: !v
          , _cpriority :: p
          , _cleft :: CTree k v p
          , _cright :: CTree k v p }
  deriving (Eq)

data Treap k v p
  = TNil
  | TNode { _tkey :: !k
          , _tvalue :: !v
          , _tpriority :: !p
          , _tleft :: Treap k v p
          , _tright :: Treap k v p }
  deriving (Eq, Show)

makeLenses ''CTree

makeLenses ''Treap

instance CartesianTree CTree k v p where
  priority = cpriority
  key = ckey
  value = cvalue
  left = cleft
  right = cright
  nil = pure CNil
  createNode k v p l r = pure (CNode k v p l r)

instance (Show k, Show v, Show p) => Show (CTree k v p) where
  show CNil = "CNil"
  show (CNode k v p l r) =
    "CNode {" ++
    "_ckey = " ++
    show k ++
    ", _cvalue = " ++
    show v ++ ", _cleft = {" ++ show l ++ "}, _cright = {" ++ show r ++ "}}"

instance CartesianTree Treap k v p where
  priority = tpriority
  key = tkey
  value = tvalue
  left = tleft
  right = tright
  nil = pure TNil
  createNode k v p l r = pure (TNode k v p l r)

data Crumb c k v p
  = LeftCrumb (Maybe k)
              (Maybe v)
              (Maybe p)
              (Maybe (c k v p))
  | RightCrumb (Maybe k)
               (Maybe v)
               (Maybe p)
               (Maybe (c k v p))

type BreadCrumbs c k v p = [Crumb c k v p]

goLeft ::
     CartesianTree c k v p
  => (c k v p, [Crumb c k v p])
  -> (Maybe (c k v p), [Crumb c k v p])
goLeft (t, crumbs) =
  ( preview left t
  , LeftCrumb
      (preview key t)
      (preview value t)
      (preview priority t)
      (preview right t) :
    crumbs)

goRight ::
     CartesianTree c k v p
  => (c k v p, [Crumb c k v p])
  -> (Maybe (c k v p), [Crumb c k v p])
goRight (t, crumbs) =
  ( preview right t
  , RightCrumb
      (preview key t)
      (preview value t)
      (preview priority t)
      (preview left t) :
    crumbs)

goUp ::
     CartesianTree c k v p
  => (c k v p, [Crumb c k v p])
  -> (c k v p, [Crumb c k v p])
goUp (t, LeftCrumb k v p r:bs) =
  let t1 = set left t t
      t2 = set value (fromJust v) t1
      t3 = set key (fromJust k) t2
      t4 = set priority (fromJust p) t3
      t5 = set right (fromJust r) t4
  in (t5, bs)
goUp (t, RightCrumb k v p l:bs) =
  let t1 = set right t t
      t2 = set value (fromJust v) t1
      t3 = set key (fromJust k) t2
      t4 = set priority (fromJust p) t3
      t5 = set left (fromJust l) t4
  in (t5, bs)

type Zipper c b a p = (c b a p, BreadCrumbs c b a p)

modifyValue ::
     (Eq (c k a p), CartesianTree c k a p)
  => (a -> a)
  -> (c k a p, b)
  -> (c k a p, b)
modifyValue f (t, bs) =
  if t == (runIdentity nil)
    then (t, bs)
    else (over value f t, bs)

modifyKey ::
     (Eq (c a v p), CartesianTree c a v p)
  => (a -> a)
  -> (c a v p, b)
  -> (c a v p, b)
modifyKey f (t, bs) =
  if t == runIdentity nil
    then (t, bs)
    else (over key f t, bs)

modifyPriority ::
     (Eq (c k v a), CartesianTree c k v a)
  => (a -> a)
  -> (c k v a, b)
  -> (c k v a, b)
modifyPriority f (t, bs) =
  if t == (runIdentity nil)
    then (t, bs)
    else (over priority f t, bs)

buildtree ::
     (Eq (c k v p), CartesianTree c k v p, Foldable t, Ord v)
  => t (k, v, p)
  -> c k v p
buildtree xs = fst $ topMost $ foldl (flip cinsert) (runIdentity nil, []) xs

cinsert ::
     (Eq (c k v p), Ord v, CartesianTree c k v p)
  => (k, v, p)
  -> (c k v p, [Crumb c k v p])
  -> (c k v p, [Crumb c k v p])
cinsert ins@(k, v, p) z@(t, bs) =
  if t == runIdentity nil
    then ( runIdentity $ createNode k v p (runIdentity nil) (runIdentity nil)
         , bs)
    else case fmap (compare v) (preview value t) of
           Just LT ->
             case bs of
               [] -> (runIdentity $ createNode k v p t (runIdentity nil), bs)
               _ -> cinsert ins (goUp z)
           Just _ ->
             let rr =
                   (runIdentity $
                    createNode
                      (fromJust $ preview key t)
                      (fromJust $ preview value t)
                      (fromJust $ preview priority t)
                      (fromJust $ preview left t)
                      (runIdentity $
                       createNode
                         k
                         v
                         p
                         (fromJust $ preview right t)
                         (runIdentity nil)))
                 (a, b) = goRight (rr, bs)
             in (fromJust a, b)

topMost ::
     CartesianTree c k v p
  => (c k v p, [Crumb c k v p])
  -> (c k v p, [Crumb c k v p])
topMost (t, []) = (t, [])
topMost z = topMost (goUp z)

testTsequence :: IO [(Int,Int,Int)]
testTsequence =
  sortBy (comparing (\(_, _, p) -> p)) <$>
  zipWithM go [1 ..] (map read $ words "9 3 7 1 8 12 10 20 15 18 5")
  where
    go k v = do
      coin <- randomRIO (0 :: Int, 1)
      priorityRand <- randomRIO (1, 1000000000 :: Int)
      return (k, v, priorityRand)

testCsequence :: [(Int, Int, Int)]
testCsequence =
  zip3 [1 ..] (map read $ words "9 3 7 1 8 12 10 20 15 18 5") (repeat undefined)

generateSequence
  :: (Eq (c k a p), CartesianTree c k a p) => c k a p -> [a]
generateSequence t =
  if t == runIdentity nil
    then []
    else generateSequence (fromJust $ preview left t) ++
         [fromJust $ preview value t] ++
         generateSequence (fromJust $ preview right t)

buildTreap :: (Ord k, Ord v, Ord p) => [(k,v,p)] -> Treap k v p
buildTreap = buildtree

buildCTree :: (Ord k, Ord v, Ord p) => [(k,v,p)] -> CTree k v p
buildCTree = buildtree