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