{-# LANGUAGE CPP #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.LCA.Online
-- Copyright   :  (C) 2011-2015 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  experimental
-- Portability :  portable
--
-- Provides online calculation of the the lowest common ancestor in /O(log h)/
-- by compressing the spine of a 'Path' using a skew-binary random access
-- list.
--
-- This library implements the technique described in my talk
--
-- <http://www.slideshare.net/ekmett/skewbinary-online-lowest-common-ancestor-search>
--
-- to improve the known asymptotic bounds on both online lowest common ancestor search
--
-- <http://en.wikipedia.org/wiki/Lowest_common_ancestor>
--
-- and the online level ancestor problem:
--
-- <http://en.wikipedia.org/wiki/Level_ancestor_problem>
--
-- Algorithms used here assume that the key values chosen for @k@ are
-- globally unique.
----------------------------------------------------------------------------
module Data.LCA.Online
  ( Path
  , lca
  , empty
  , cons
  , uncons, view
  , null
  , length
  , isAncestorOf
  , keep
  , drop
  , traverseWithKey
  , toList
  , fromList
  , (~=)
  ) where

import qualified Data.Foldable as F

#if __GLASGOW_HASKELL__ < 710
import Control.Applicative hiding (empty)
import Data.Traversable
import Data.Monoid
#endif

import Data.LCA.View

import Prelude hiding
  ( drop
  , length
  , null
  )

-- $setup
-- >>> let length = Data.LCA.Online.length

-- | Complete binary trees
data Tree a
  = Bin {-# UNPACK #-} !Int a (Tree a) (Tree a)
  | Tip {-# UNPACK #-} !Int a
  deriving (Int -> Tree a -> ShowS
[Tree a] -> ShowS
Tree a -> String
(Int -> Tree a -> ShowS)
-> (Tree a -> String) -> ([Tree a] -> ShowS) -> Show (Tree a)
forall a. Show a => Int -> Tree a -> ShowS
forall a. Show a => [Tree a] -> ShowS
forall a. Show a => Tree a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Tree a] -> ShowS
$cshowList :: forall a. Show a => [Tree a] -> ShowS
show :: Tree a -> String
$cshow :: forall a. Show a => Tree a -> String
showsPrec :: Int -> Tree a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Tree a -> ShowS
Show, ReadPrec [Tree a]
ReadPrec (Tree a)
Int -> ReadS (Tree a)
ReadS [Tree a]
(Int -> ReadS (Tree a))
-> ReadS [Tree a]
-> ReadPrec (Tree a)
-> ReadPrec [Tree a]
-> Read (Tree a)
forall a. Read a => ReadPrec [Tree a]
forall a. Read a => ReadPrec (Tree a)
forall a. Read a => Int -> ReadS (Tree a)
forall a. Read a => ReadS [Tree a]
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [Tree a]
$creadListPrec :: forall a. Read a => ReadPrec [Tree a]
readPrec :: ReadPrec (Tree a)
$creadPrec :: forall a. Read a => ReadPrec (Tree a)
readList :: ReadS [Tree a]
$creadList :: forall a. Read a => ReadS [Tree a]
readsPrec :: Int -> ReadS (Tree a)
$creadsPrec :: forall a. Read a => Int -> ReadS (Tree a)
Read)

instance Functor Tree where
  fmap :: (a -> b) -> Tree a -> Tree b
fmap a -> b
f (Bin Int
n a
a Tree a
l Tree a
r) = Int -> b -> Tree b -> Tree b -> Tree b
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
n (a -> b
f a
a) ((a -> b) -> Tree a -> Tree b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Tree a
l) ((a -> b) -> Tree a -> Tree b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Tree a
r)
  fmap a -> b
f (Tip Int
n a
a)     = Int -> b -> Tree b
forall a. Int -> a -> Tree a
Tip Int
n (a -> b
f a
a)

instance F.Foldable Tree where
  foldMap :: (a -> m) -> Tree a -> m
foldMap a -> m
f (Bin Int
_ a
a Tree a
l Tree a
r) = a -> m
f a
a m -> m -> m
forall a. Monoid a => a -> a -> a
`mappend` (a -> m) -> Tree a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
F.foldMap a -> m
f Tree a
l m -> m -> m
forall a. Monoid a => a -> a -> a
`mappend` (a -> m) -> Tree a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
F.foldMap a -> m
f Tree a
r
  foldMap a -> m
f (Tip Int
_ a
a)     = a -> m
f a
a

instance Traversable Tree where
  traverse :: (a -> f b) -> Tree a -> f (Tree b)
traverse a -> f b
f (Bin Int
n a
a Tree a
l Tree a
r) = Int -> b -> Tree b -> Tree b -> Tree b
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
n (b -> Tree b -> Tree b -> Tree b)
-> f b -> f (Tree b -> Tree b -> Tree b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
f a
a f (Tree b -> Tree b -> Tree b)
-> f (Tree b) -> f (Tree b -> Tree b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (a -> f b) -> Tree a -> f (Tree b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Tree a
l f (Tree b -> Tree b) -> f (Tree b) -> f (Tree b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (a -> f b) -> Tree a -> f (Tree b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Tree a
r
  traverse a -> f b
f (Tip Int
n a
a)     = Int -> b -> Tree b
forall a. Int -> a -> Tree a
Tip Int
n (b -> Tree b) -> f b -> f (Tree b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
f a
a

sameT :: Tree a -> Tree b -> Bool
sameT :: Tree a -> Tree b -> Bool
sameT Tree a
xs Tree b
ys = Tree a -> Int
forall a. Tree a -> Int
root Tree a
xs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Tree b -> Int
forall a. Tree a -> Int
root Tree b
ys where
  root :: Tree a -> Int
root (Tip Int
k a
_)     = Int
k
  root (Bin Int
k a
_ Tree a
_ Tree a
_) = Int
k
{-# INLINE sameT #-}

-- | Compressed paths using skew binary random access lists
data Path a
  = Nil
  | Cons {-# UNPACK #-} !Int -- the number of elements @n@ in this entire skew list
         {-# UNPACK #-} !Int -- the number of elements @w@ in this binary tree node
         (Tree a)            -- a complete binary tree @t@ of with @w@ elements
         (Path a)            -- @n - w@ elements in a linked list @ts@, of complete trees in ascending order by size
  deriving (Int -> Path a -> ShowS
[Path a] -> ShowS
Path a -> String
(Int -> Path a -> ShowS)
-> (Path a -> String) -> ([Path a] -> ShowS) -> Show (Path a)
forall a. Show a => Int -> Path a -> ShowS
forall a. Show a => [Path a] -> ShowS
forall a. Show a => Path a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Path a] -> ShowS
$cshowList :: forall a. Show a => [Path a] -> ShowS
show :: Path a -> String
$cshow :: forall a. Show a => Path a -> String
showsPrec :: Int -> Path a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Path a -> ShowS
Show, ReadPrec [Path a]
ReadPrec (Path a)
Int -> ReadS (Path a)
ReadS [Path a]
(Int -> ReadS (Path a))
-> ReadS [Path a]
-> ReadPrec (Path a)
-> ReadPrec [Path a]
-> Read (Path a)
forall a. Read a => ReadPrec [Path a]
forall a. Read a => ReadPrec (Path a)
forall a. Read a => Int -> ReadS (Path a)
forall a. Read a => ReadS [Path a]
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [Path a]
$creadListPrec :: forall a. Read a => ReadPrec [Path a]
readPrec :: ReadPrec (Path a)
$creadPrec :: forall a. Read a => ReadPrec (Path a)
readList :: ReadS [Path a]
$creadList :: forall a. Read a => ReadS [Path a]
readsPrec :: Int -> ReadS (Path a)
$creadsPrec :: forall a. Read a => Int -> ReadS (Path a)
Read)

instance Functor Path where
  fmap :: (a -> b) -> Path a -> Path b
fmap a -> b
_ Path a
Nil = Path b
forall a. Path a
Nil
  fmap a -> b
f (Cons Int
n Int
k Tree a
t Path a
ts) = Int -> Int -> Tree b -> Path b -> Path b
forall a. Int -> Int -> Tree a -> Path a -> Path a
Cons Int
n Int
k ((a -> b) -> Tree a -> Tree b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Tree a
t) ((a -> b) -> Path a -> Path b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Path a
ts)

instance F.Foldable Path where
  foldMap :: (a -> m) -> Path a -> m
foldMap a -> m
_ Path a
Nil = m
forall a. Monoid a => a
mempty
  foldMap a -> m
f (Cons Int
_ Int
_ Tree a
t Path a
ts) = (a -> m) -> Tree a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
F.foldMap a -> m
f Tree a
t m -> m -> m
forall a. Monoid a => a -> a -> a
`mappend` (a -> m) -> Path a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
F.foldMap a -> m
f Path a
ts

#if __GLASGOW_HASKELL__ >= 710
  length :: Path a -> Int
length = Path a -> Int
forall a. Path a -> Int
length
  null :: Path a -> Bool
null   = Path a -> Bool
forall a. Path a -> Bool
null
#endif

-- | /O(1)/ Determine the 'length' of a 'Path'.
length :: Path a -> Int
length :: Path a -> Int
length Path a
Nil = Int
0
length (Cons Int
n Int
_ Tree a
_ Path a
_) = Int
n
{-# INLINE length #-}

-- | /O(1)/ Returns 'True' iff the path is 'empty'.
null :: Path a -> Bool
null :: Path a -> Bool
null Path a
Nil = Bool
True
null Path a
_ = Bool
False
{-# INLINE null #-}

instance Traversable Path where
  traverse :: (a -> f b) -> Path a -> f (Path b)
traverse a -> f b
_ Path a
Nil = Path b -> f (Path b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Path b
forall a. Path a
Nil
  traverse a -> f b
f (Cons Int
n Int
k Tree a
t Path a
ts) = Int -> Int -> Tree b -> Path b -> Path b
forall a. Int -> Int -> Tree a -> Path a -> Path a
Cons Int
n Int
k (Tree b -> Path b -> Path b) -> f (Tree b) -> f (Path b -> Path b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> Tree a -> f (Tree b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Tree a
t f (Path b -> Path b) -> f (Path b) -> f (Path b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (a -> f b) -> Path a -> f (Path b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Path a
ts

consT :: Int -> Tree a -> Path a -> Path a
consT :: Int -> Tree a -> Path a -> Path a
consT Int
w Tree a
t Path a
ts = Int -> Int -> Tree a -> Path a -> Path a
forall a. Int -> Int -> Tree a -> Path a -> Path a
Cons (Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Path a -> Int
forall a. Path a -> Int
length Path a
ts) Int
w Tree a
t Path a
ts
{-# INLINE consT #-}

-- | Convert a 'Path' to a list of @(ID, value)@ pairs.
toList :: Path a -> [(Int,a)]
toList :: Path a -> [(Int, a)]
toList Path a
Nil = []
toList (Cons Int
_ Int
_ Tree a
t Path a
ts) = Tree a -> [(Int, a)] -> [(Int, a)]
forall b. Tree b -> [(Int, b)] -> [(Int, b)]
go Tree a
t (Path a -> [(Int, a)]
forall a. Path a -> [(Int, a)]
toList Path a
ts) where
  go :: Tree b -> [(Int, b)] -> [(Int, b)]
go (Tip Int
k b
a) [(Int, b)]
xs     = (Int
k,b
a) (Int, b) -> [(Int, b)] -> [(Int, b)]
forall a. a -> [a] -> [a]
: [(Int, b)]
xs
  go (Bin Int
k b
a Tree b
l Tree b
r) [(Int, b)]
xs = (Int
k,b
a) (Int, b) -> [(Int, b)] -> [(Int, b)]
forall a. a -> [a] -> [a]
: Tree b -> [(Int, b)] -> [(Int, b)]
go Tree b
l (Tree b -> [(Int, b)] -> [(Int, b)]
go Tree b
r [(Int, b)]
xs)

-- | Build a 'Path' from a list of @(ID, value)@ pairs.
fromList :: [(Int,a)] -> Path a
fromList :: [(Int, a)] -> Path a
fromList [] = Path a
forall a. Path a
Nil
fromList ((Int
k,a
a):[(Int, a)]
xs) = Int -> a -> Path a -> Path a
forall a. Int -> a -> Path a -> Path a
cons Int
k a
a ([(Int, a)] -> Path a
forall a. [(Int, a)] -> Path a
fromList [(Int, a)]
xs)

-- | Traverse a 'Path' with access to the node IDs.
traverseWithKey :: Applicative f => (Int -> a -> f b) -> Path a -> f (Path b)
traverseWithKey :: (Int -> a -> f b) -> Path a -> f (Path b)
traverseWithKey Int -> a -> f b
f = Path a -> f (Path b)
go where
  go :: Path a -> f (Path b)
go Path a
Nil = Path b -> f (Path b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Path b
forall a. Path a
Nil
  go (Cons Int
n Int
k Tree a
t Path a
ts) = Int -> Int -> Tree b -> Path b -> Path b
forall a. Int -> Int -> Tree a -> Path a -> Path a
Cons Int
n Int
k (Tree b -> Path b -> Path b) -> f (Tree b) -> f (Path b -> Path b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Tree a -> f (Tree b)
goT Tree a
t f (Path b -> Path b) -> f (Path b) -> f (Path b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Path a -> f (Path b)
go Path a
ts
  goT :: Tree a -> f (Tree b)
goT (Bin Int
k a
a Tree a
l Tree a
r) = Int -> b -> Tree b -> Tree b -> Tree b
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
k (b -> Tree b -> Tree b -> Tree b)
-> f b -> f (Tree b -> Tree b -> Tree b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> a -> f b
f Int
k a
a f (Tree b -> Tree b -> Tree b)
-> f (Tree b) -> f (Tree b -> Tree b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Tree a -> f (Tree b)
goT Tree a
l f (Tree b -> Tree b) -> f (Tree b) -> f (Tree b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Tree a -> f (Tree b)
goT Tree a
r
  goT (Tip Int
k a
a)     = Int -> b -> Tree b
forall a. Int -> a -> Tree a
Tip Int
k (b -> Tree b) -> f b -> f (Tree b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> a -> f b
f Int
k a
a
{-# INLINE traverseWithKey #-}

-- | The 'empty' 'Path'
empty :: Path a
empty :: Path a
empty = Path a
forall a. Path a
Nil
{-# INLINE empty #-}

-- | /O(1)/ Invariant: most operations assume that the keys @k@ are globally unique
--
-- Extend the 'Path' with a new node ID and value.
cons :: Int -> a -> Path a -> Path a
cons :: Int -> a -> Path a -> Path a
cons Int
k a
a (Cons Int
n Int
w Tree a
t (Cons Int
_ Int
w' Tree a
t2 Path a
ts)) | Int
w Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
w' = Int -> Int -> Tree a -> Path a -> Path a
forall a. Int -> Int -> Tree a -> Path a -> Path a
Cons (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
k a
a Tree a
t Tree a
t2) Path a
ts
cons Int
k a
a Path a
ts = Int -> Int -> Tree a -> Path a -> Path a
forall a. Int -> Int -> Tree a -> Path a -> Path a
Cons (Path a -> Int
forall a. Path a -> Int
length Path a
ts Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
1 (Int -> a -> Tree a
forall a. Int -> a -> Tree a
Tip Int
k a
a) Path a
ts
{-# INLINE cons #-}

-- | /O(1)/ Extract the node ID and value from the newest node on the 'Path'.
uncons :: Path a -> Maybe (Int, a, Path a)
uncons :: Path a -> Maybe (Int, a, Path a)
uncons Path a
Nil = Maybe (Int, a, Path a)
forall a. Maybe a
Nothing
uncons (Cons Int
_ Int
_ (Tip Int
k a
a) Path a
ts) = (Int, a, Path a) -> Maybe (Int, a, Path a)
forall a. a -> Maybe a
Just (Int
k, a
a, Path a
ts)
uncons (Cons Int
_ Int
w (Bin Int
k a
a Tree a
l Tree a
r) Path a
ts) = (Int, a, Path a) -> Maybe (Int, a, Path a)
forall a. a -> Maybe a
Just (Int
k, a
a, Int -> Tree a -> Path a -> Path a
forall a. Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
l (Int -> Tree a -> Path a -> Path a
forall a. Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
r Path a
ts)) where w2 :: Int
w2 = Int -> Int -> Int
forall a. Integral a => a -> a -> a
div Int
w Int
2
{-# INLINE uncons #-}

-- | /O(1)/ Extract the node ID and value from the newest node on the 'Path', slightly faster than 'uncons'.
view :: Path a -> View Path a
view :: Path a -> View Path a
view Path a
Nil = View Path a
forall (f :: * -> *) a. View f a
Root
view (Cons Int
_ Int
_ (Tip Int
k a
a) Path a
ts) = Int -> a -> Path a -> View Path a
forall (f :: * -> *) a. Int -> a -> f a -> View f a
Node Int
k a
a Path a
ts
view (Cons Int
_ Int
w (Bin Int
k a
a Tree a
l Tree a
r) Path a
ts) = Int -> a -> Path a -> View Path a
forall (f :: * -> *) a. Int -> a -> f a -> View f a
Node Int
k a
a (Int -> Tree a -> Path a -> Path a
forall a. Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
l (Int -> Tree a -> Path a -> Path a
forall a. Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
r Path a
ts)) where w2 :: Int
w2 = Int -> Int -> Int
forall a. Integral a => a -> a -> a
div Int
w Int
2
{-# INLINE view #-}

-- | /O(log (h - k))/ to @'keep' k@ elements of 'Path' of 'length' @h@
--
-- This solves the online version of the \"level ancestor problem\" with no preprocessing in /O(log h)/ time,
-- improving known complexity bounds.
--
-- <http://en.wikipedia.org/wiki/Level_ancestor_problem>
keep :: Int -> Path a -> Path a
keep :: Int -> Path a -> Path a
keep = Int -> Path a -> Path a
forall a. Int -> Path a -> Path a
go where
  go :: Int -> Path a -> Path a
go Int
_ Path a
Nil = Path a
forall a. Path a
Nil
  go Int
k xs :: Path a
xs@(Cons Int
n Int
w Tree a
t Path a
ts)
    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n    = Path a
xs
    | Bool
otherwise = case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
k (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
w) of
      Ordering
GT -> Int -> Int -> Tree a -> Path a -> Path a
forall a. Int -> Int -> Tree a -> Path a -> Path a
goT (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
w) Int
w Tree a
t Path a
ts
      Ordering
EQ -> Path a
ts
      Ordering
LT -> Int -> Path a -> Path a
go Int
k Path a
ts
  goT :: Int -> Int -> Tree a -> Path a -> Path a
goT Int
n Int
w (Bin Int
_ a
_ Tree a
l Tree a
r) Path a
ts = case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
n Int
w2 of
    Ordering
LT              -> Int -> Int -> Tree a -> Path a -> Path a
goT Int
n Int
w2 Tree a
r Path a
ts
    Ordering
EQ              -> Int -> Tree a -> Path a -> Path a
forall a. Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
r Path a
ts
    Ordering
GT | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 -> Int -> Tree a -> Path a -> Path a
forall a. Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
l (Int -> Tree a -> Path a -> Path a
forall a. Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
r Path a
ts)
       | Bool
otherwise  -> Int -> Int -> Tree a -> Path a -> Path a
goT (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
w2) Int
w2 Tree a
l (Int -> Tree a -> Path a -> Path a
forall a. Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
r Path a
ts)
    where w2 :: Int
w2 = Int -> Int -> Int
forall a. Integral a => a -> a -> a
div Int
w Int
2
  goT Int
_ Int
_ Tree a
_ Path a
ts = Path a
ts
{-# INLINE keep #-}

-- | /O(log k)/ to @'drop' k@ elements from a 'Path'
drop :: Int -> Path a -> Path a
drop :: Int -> Path a -> Path a
drop Int
k Path a
xs = Int -> Path a -> Path a
forall a. Int -> Path a -> Path a
keep (Path a -> Int
forall a. Path a -> Int
length Path a
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k) Path a
xs
{-# INLINE drop #-}

-- | /O(log h)/ @xs `isAncestorOf` ys@ holds when @xs@ is a prefix starting at the root of 'Path' @ys@.
isAncestorOf :: Path a -> Path b -> Bool
isAncestorOf :: Path a -> Path b -> Bool
isAncestorOf Path a
xs Path b
ys = Path a
xs Path a -> Path b -> Bool
forall a b. Path a -> Path b -> Bool
~= Int -> Path b -> Path b
forall a. Int -> Path a -> Path a
keep (Path a -> Int
forall a. Path a -> Int
length Path a
xs) Path b
ys
{-# INLINE isAncestorOf #-}

infix 4 ~=
-- | /O(1)/ Compare to see if two trees have the same root key
(~=) :: Path a -> Path b -> Bool
Path a
Nil          ~= :: Path a -> Path b -> Bool
~= Path b
Nil          = Bool
True
Cons Int
_ Int
_ Tree a
s Path a
_ ~= Cons Int
_ Int
_ Tree b
t Path b
_ = Tree a -> Tree b -> Bool
forall a b. Tree a -> Tree b -> Bool
sameT Tree a
s Tree b
t
Path a
_            ~= Path b
_            = Bool
False
{-# INLINE (~=) #-}

-- | /O(log h)/ Compute the lowest common ancestor of two paths.
--
-- >>> let fromList' = fromList . map (flip (,) ())
-- >>> length (lca (fromList' [1, 2, 3, 4, 5, 6]) (fromList' [7, 8, 3, 4, 5, 6]))
-- 4
--
lca :: Path a -> Path b -> Path a
lca :: Path a -> Path b -> Path a
lca Path a
xs0 Path b
ys0 = case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
nxs Int
nys of
    Ordering
LT -> Path a -> Path b -> Path a
forall a b. Path a -> Path b -> Path a
go Path a
xs0 (Int -> Path b -> Path b
forall a. Int -> Path a -> Path a
keep Int
nxs Path b
ys0)
    Ordering
EQ -> Path a -> Path b -> Path a
forall a b. Path a -> Path b -> Path a
go Path a
xs0 Path b
ys0
    Ordering
GT -> Path a -> Path b -> Path a
forall a b. Path a -> Path b -> Path a
go (Int -> Path a -> Path a
forall a. Int -> Path a -> Path a
keep Int
nys Path a
xs0) Path b
ys0
  where
    nxs :: Int
nxs = Path a -> Int
forall a. Path a -> Int
length Path a
xs0
    nys :: Int
nys = Path b -> Int
forall a. Path a -> Int
length Path b
ys0
    go :: Path a -> Path b -> Path a
go h :: Path a
h@(Cons Int
_ Int
w Tree a
x Path a
xs) (Cons Int
_ Int
_ Tree b
y Path b
ys)
      | Tree a -> Tree b -> Bool
forall a b. Tree a -> Tree b -> Bool
sameT Tree a
x Tree b
y = Path a
h
      | Path a
xs Path a -> Path b -> Bool
forall a b. Path a -> Path b -> Bool
~= Path b
ys  = Int -> Tree a -> Tree b -> Path a -> Path a
forall a b. Int -> Tree a -> Tree b -> Path a -> Path a
goT Int
w Tree a
x Tree b
y Path a
xs
      | Bool
otherwise = Path a -> Path b -> Path a
go Path a
xs Path b
ys
    go Path a
_ Path b
_ = Path a
forall a. Path a
Nil

    goT :: Int -> Tree a -> Tree b -> Path a -> Path a
goT Int
w (Bin Int
_ a
_ Tree a
la Tree a
ra) (Bin Int
_ b
_ Tree b
lb Tree b
rb) Path a
ts
      | Tree a -> Tree b -> Bool
forall a b. Tree a -> Tree b -> Bool
sameT Tree a
la Tree b
lb = Int -> Tree a -> Path a -> Path a
forall a. Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
la (Int -> Tree a -> Path a -> Path a
forall a. Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
ra Path a
ts)
      | Tree a -> Tree b -> Bool
forall a b. Tree a -> Tree b -> Bool
sameT Tree a
ra Tree b
rb = Int -> Tree a -> Tree b -> Path a -> Path a
goT Int
w2 Tree a
la Tree b
lb (Int -> Tree a -> Path a -> Path a
forall a. Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
ra Path a
ts)
      | Bool
otherwise   = Int -> Tree a -> Tree b -> Path a -> Path a
goT Int
w2 Tree a
ra Tree b
rb Path a
ts
      where w2 :: Int
w2 = Int -> Int -> Int
forall a. Integral a => a -> a -> a
div Int
w Int
2
    goT Int
_ Tree a
_ Tree b
_ Path a
ts = Path a
ts
{-# INLINE lca #-}