module Data.LCA.Online
( Path
, empty
, cons
, null
, length
, isAncestorOf
, lca
, keep
, drop
, traverseWithKey
, toList
, fromList
, (~=)
) where
import Control.Applicative hiding (empty)
import Data.Foldable hiding (toList)
import Data.Traversable
import Data.Monoid
import Prelude hiding (length, null, drop)
data Path k a
= Nil
| Cons !Int
!Int
(Tree k a)
(Path k a)
deriving (Show, Read)
instance Functor (Path k) where
fmap _ Nil = Nil
fmap f (Cons n k t ts) = Cons n k (fmap f t) (fmap f ts)
instance Foldable (Path k) where
foldMap _ Nil = mempty
foldMap f (Cons _ _ t ts) = foldMap f t `mappend` foldMap f ts
instance Traversable (Path k) where
traverse _ Nil = pure Nil
traverse f (Cons n k t ts) = Cons n k <$> traverse f t <*> traverse f ts
data Tree k a
= Bin k a (Tree k a) (Tree k a)
| Tip k a
deriving (Show, Read)
instance Functor (Tree k) where
fmap f (Bin n a l r) = Bin n (f a) (fmap f l) (fmap f r)
fmap f (Tip n a) = Tip n (f a)
instance Foldable (Tree k) where
foldMap f (Bin _ a l r) = f a `mappend` foldMap f l `mappend` foldMap f r
foldMap f (Tip _ a) = f a
instance Traversable (Tree k) where
traverse f (Bin n a l r) = Bin n <$> f a <*> traverse f l <*> traverse f r
traverse f (Tip n a) = Tip n <$> f a
toList :: Path k a -> [(k,a)]
toList Nil = []
toList (Cons _ _ t ts) = go t (toList ts) where
go (Tip k a) xs = (k,a) : xs
go (Bin k a l r) xs = (k,a) : go l (go r xs)
fromList :: [(k,a)] -> Path k a
fromList [] = Nil
fromList ((k,a):xs) = cons k a (fromList xs)
traverseWithKey :: Applicative f => (k -> a -> f b) -> Path k a -> f (Path k b)
traverseWithKey _ Nil = pure Nil
traverseWithKey f (Cons n k t ts) = Cons n k <$> traverseTreeWithKey f t <*> traverseWithKey f ts
empty :: Path k a
empty = Nil
length :: Path k a -> Int
length Nil = 0
length (Cons n _ _ _) = n
null :: Path k a -> Bool
null Nil = True
null _ = False
cons :: k -> a -> Path k a -> Path k a
cons k a (Cons n w t (Cons _ w' t2 ts)) | w == w' = Cons (n + 1) (2 * w + 1) (Bin k a t t2) ts
cons k a ts = Cons (length ts + 1) 1 (Tip k a) ts
keep :: Int -> Path k a -> Path k a
keep _ Nil = Nil
keep k xs@(Cons n w t ts)
| k >= n = xs
| otherwise = case compare k (n w) of
GT -> keepT (k n + w) w t ts
EQ -> ts
LT -> keep k ts
drop :: Int -> Path k a -> Path k a
drop k xs = keep (length xs k) xs
lca :: Eq k => Path k a -> Path k b -> Path k a
lca xs ys = case compare nxs nys of
LT -> lca' xs (keep nxs ys)
EQ -> lca' xs ys
GT -> lca' (keep nys xs) ys
where
nxs = length xs
nys = length ys
isAncestorOf :: Eq k => Path k a -> Path k b -> Bool
isAncestorOf xs ys = xs ~= keep (length xs) ys
infix 4 ~=
(~=) :: Eq k => Path k a -> Path k b -> Bool
Nil ~= Nil = True
Cons _ _ s _ ~= Cons _ _ t _ = sameT s t
_ ~= _ = False
consT :: Int -> Tree k a -> Path k a -> Path k a
consT w t ts = Cons (w + length ts) w t ts
keepT :: Int -> Int -> Tree k a -> Path k a -> Path k a
keepT n w (Bin _ _ l r) ts = case compare n w2 of
LT -> keepT n w2 r ts
EQ -> consT w2 r ts
GT | n == w 1 -> consT w2 l (consT w2 r ts)
| otherwise -> keepT (n w2) w2 l (consT w2 r ts)
where w2 = div w 2
keepT _ _ _ ts = ts
sameT :: Eq k => Tree k a -> Tree k b -> Bool
sameT xs ys = root xs == root ys
lca' :: Eq k => Path k a -> Path k b -> Path k a
lca' h@(Cons _ w x xs) (Cons _ _ y ys)
| sameT x y = h
| xs ~= ys = lcaT w x y xs
| otherwise = lca' xs ys
lca' _ _ = Nil
lcaT :: Eq k => Int -> Tree k a -> Tree k b -> Path k a -> Path k a
lcaT w (Bin _ _ la ra) (Bin _ _ lb rb) ts
| sameT la lb = consT w2 la (consT w2 ra ts)
| sameT ra rb = lcaT w2 la lb (consT w ra ts)
| otherwise = lcaT w2 ra rb ts
where w2 = div w 2
lcaT _ _ _ ts = ts
traverseTreeWithKey :: Applicative f => (k -> a -> f b) -> Tree k a -> f (Tree k b)
traverseTreeWithKey f (Bin k a l r) = Bin k <$> f k a <*> traverseTreeWithKey f l <*> traverseTreeWithKey f r
traverseTreeWithKey f (Tip k a) = Tip k <$> f k a
root :: Tree k a -> k
root (Tip k _) = k
root (Bin k _ _ _) = k