{-# LANGUAGE ScopedTypeVariables #-}
{-|
== Prufer sequences

Functions to convert trees to Prufer sequences and vice versa

Sources:

* https://en.wikipedia.org/wiki/Pr%C3%BCfer_sequence
* Xiaodong Wang, Lei Wang, Yingjie Wu, "An Optimal Algorithm for Prufer Codes", 2009
  https://www.scirp.org/pdf/JSEA20090200006_93737200.pdf

-}

{-
Implementation notes:
* The c array tracks the number of children of a node.
-}

module Prufer
    ( graphToSeq
    , treeToSeq
    , seqToEdges
    , seqToGraph
    ) where

import Control.Monad
import Control.Monad.ST
import Data.Array.Unboxed
import Data.Array.ST
import Data.Graph

nxtLeaf :: STUArray s Vertex Int -> Vertex -> ST s Vertex
nxtLeaf :: STUArray s Vertex Vertex -> Vertex -> ST s Vertex
nxtLeaf STUArray s Vertex Vertex
c Vertex
u = do
    Vertex
cu <- STUArray s Vertex Vertex -> Vertex -> ST s Vertex
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Vertex Vertex
c Vertex
u
    if Vertex
cu Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
0 then Vertex -> ST s Vertex
forall (m :: * -> *) a. Monad m => a -> m a
return Vertex
u else STUArray s Vertex Vertex -> Vertex -> ST s Vertex
forall s. STUArray s Vertex Vertex -> Vertex -> ST s Vertex
nxtLeaf STUArray s Vertex Vertex
c (Vertex -> ST s Vertex) -> Vertex -> ST s Vertex
forall a b. (a -> b) -> a -> b
$ Vertex
u Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
1

withSizeCheck :: Vertex -> Vertex -> [a] -> [a]
withSizeCheck :: Vertex -> Vertex -> [a] -> [a]
withSizeCheck Vertex
l Vertex
r [a]
xs
    | Vertex
l Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
> Vertex
r Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
1 = [Char] -> [a]
forall a. HasCallStack => [Char] -> a
error [Char]
"negative node count"
    | Vertex
l Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
>= Vertex
r    = []
    | Bool
otherwise = [a]
xs

-- | Convert a bidirected Graph to a Prufer sequence. The graph must be a connected tree or empty. O(n).
graphToSeq :: Graph -> [Vertex]
graphToSeq :: Graph -> [Vertex]
graphToSeq Graph
g = Bounds -> Tree Vertex -> [Vertex]
treeToSeq Bounds
b Tree Vertex
t where
    b :: Bounds
b@(Vertex
_, Vertex
r) = Graph -> Bounds
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
bounds Graph
g
    [Tree Vertex
t] = Graph -> [Vertex] -> [Tree Vertex]
dfs Graph
g [Vertex
r]

-- | Convert a Tree to a Prufer sequence. The tree must be rooted at r. O(n).
treeToSeq :: Bounds -> Tree Vertex -> [Vertex]
treeToSeq :: Bounds -> Tree Vertex -> [Vertex]
treeToSeq (Vertex
l, Vertex
r) Tree Vertex
t = Vertex -> Vertex -> [Vertex] -> [Vertex]
forall a. Vertex -> Vertex -> [a] -> [a]
withSizeCheck Vertex
l Vertex
r ([Vertex] -> [Vertex]) -> [Vertex] -> [Vertex]
forall a b. (a -> b) -> a -> b
$ (forall s. ST s [Vertex]) -> [Vertex]
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s [Vertex]) -> [Vertex])
-> (forall s. ST s [Vertex]) -> [Vertex]
forall a b. (a -> b) -> a -> b
$ do
    STUArray s Vertex Vertex
c :: STUArray s Vertex Int <- Bounds -> Vertex -> ST s (STUArray s Vertex Vertex)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Vertex
l, Vertex
r) Vertex
0
    let setC :: Tree Vertex -> ST s ()
setC (Node Vertex
x [Tree Vertex]
ts) = STUArray s Vertex Vertex -> Vertex -> Vertex -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Vertex Vertex
c Vertex
x ([Tree Vertex] -> Vertex
forall (t :: * -> *) a. Foldable t => t a -> Vertex
length [Tree Vertex]
ts) ST s () -> ST s () -> ST s ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (Tree Vertex -> ST s ()) -> [Tree Vertex] -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Tree Vertex -> ST s ()
setC [Tree Vertex]
ts :: ST s ()
        par :: UArray Vertex Vertex
par = Bounds -> [Bounds] -> UArray Vertex Vertex
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [(i, e)] -> a i e
array (Vertex
l, Vertex
r) ([Bounds] -> UArray Vertex Vertex)
-> [Bounds] -> UArray Vertex Vertex
forall a b. (a -> b) -> a -> b
$ Vertex -> Tree Vertex -> [Bounds] -> [Bounds]
forall t. t -> Tree t -> [(t, t)] -> [(t, t)]
go Vertex
r Tree Vertex
t [] :: UArray Vertex Vertex where
            go :: t -> Tree t -> [(t, t)] -> [(t, t)]
go t
p (Node t
u Forest t
ts) [(t, t)]
acc = (t
u, t
p) (t, t) -> [(t, t)] -> [(t, t)]
forall a. a -> [a] -> [a]
: (Tree t -> [(t, t)] -> [(t, t)])
-> [(t, t)] -> Forest t -> [(t, t)]
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (t -> Tree t -> [(t, t)] -> [(t, t)]
go t
u) [(t, t)]
acc Forest t
ts
        go :: Vertex -> Vertex -> ST s [Vertex]
go Vertex
k Vertex
u = do
            let p :: Vertex
p = UArray Vertex Vertex
parUArray Vertex Vertex -> Vertex -> Vertex
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Vertex
u
            Vertex
cp <- STUArray s Vertex Vertex -> Vertex -> ST s Vertex
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Vertex Vertex
c Vertex
p
            if Vertex
p Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
r Bool -> Bool -> Bool
&& Vertex
cp Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
1 then [Vertex] -> ST s [Vertex]
forall (m :: * -> *) a. Monad m => a -> m a
return [] else do
                STUArray s Vertex Vertex -> Vertex -> Vertex -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Vertex Vertex
c Vertex
p (Vertex -> ST s ()) -> Vertex -> ST s ()
forall a b. (a -> b) -> a -> b
$ Vertex
cp Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- Vertex
1
                (Vertex
pVertex -> [Vertex] -> [Vertex]
forall a. a -> [a] -> [a]
:) ([Vertex] -> [Vertex]) -> ST s [Vertex] -> ST s [Vertex]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> if Vertex
p Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
< Vertex
k Bool -> Bool -> Bool
&& Vertex
cp Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
1
                            then Vertex -> Vertex -> ST s [Vertex]
go Vertex
k Vertex
p
                            else STUArray s Vertex Vertex -> Vertex -> ST s Vertex
forall s. STUArray s Vertex Vertex -> Vertex -> ST s Vertex
nxtLeaf STUArray s Vertex Vertex
c (Vertex
k Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
1) ST s Vertex -> (Vertex -> ST s [Vertex]) -> ST s [Vertex]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Vertex -> Vertex -> ST s [Vertex]) -> Vertex -> ST s [Vertex]
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join Vertex -> Vertex -> ST s [Vertex]
go
    Tree Vertex -> ST s ()
setC Tree Vertex
t ST s () -> ST s Vertex -> ST s Vertex
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> STUArray s Vertex Vertex -> Vertex -> ST s Vertex
forall s. STUArray s Vertex Vertex -> Vertex -> ST s Vertex
nxtLeaf STUArray s Vertex Vertex
c Vertex
l ST s Vertex -> (Vertex -> ST s [Vertex]) -> ST s [Vertex]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Vertex -> Vertex -> ST s [Vertex]) -> Vertex -> ST s [Vertex]
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join Vertex -> Vertex -> ST s [Vertex]
go

-- | Convert a Prufer sequence to a bidirected Graph. The sequence must be valid. O(n).
seqToGraph :: Bounds -> [Vertex] -> Graph
seqToGraph :: Bounds -> [Vertex] -> Graph
seqToGraph Bounds
bnds = Bounds -> [Bounds] -> Graph
buildG Bounds
bnds ([Bounds] -> Graph) -> ([Vertex] -> [Bounds]) -> [Vertex] -> Graph
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bounds -> [Vertex] -> [Bounds]
seqToEdges Bounds
bnds

-- | Convert a Prufer sequence to a list of edges, each edge present twice in its two orientations. The
-- sequence must be valid. O(n).
seqToEdges :: Bounds -> [Vertex] -> [Edge]
seqToEdges :: Bounds -> [Vertex] -> [Bounds]
seqToEdges (Vertex
l, Vertex
r) [Vertex]
ps = Vertex -> Vertex -> [Bounds] -> [Bounds]
forall a. Vertex -> Vertex -> [a] -> [a]
withSizeCheck Vertex
l Vertex
r ([Bounds] -> [Bounds]) -> [Bounds] -> [Bounds]
forall a b. (a -> b) -> a -> b
$ (forall s. ST s [Bounds]) -> [Bounds]
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s [Bounds]) -> [Bounds])
-> (forall s. ST s [Bounds]) -> [Bounds]
forall a b. (a -> b) -> a -> b
$ do
    STUArray s Vertex Vertex
c :: STUArray s Vertex Int <- Bounds -> Vertex -> ST s (STUArray s Vertex Vertex)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Vertex
l, Vertex
r) Vertex
0
    [Vertex] -> (Vertex -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Vertex]
ps ((Vertex -> ST s ()) -> ST s ()) -> (Vertex -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Vertex
u -> STUArray s Vertex Vertex -> Vertex -> Vertex -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Vertex Vertex
c Vertex
u (Vertex -> ST s ()) -> (Vertex -> Vertex) -> Vertex -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+Vertex
1) (Vertex -> ST s ()) -> ST s Vertex -> ST s ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< STUArray s Vertex Vertex -> Vertex -> ST s Vertex
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Vertex Vertex
c Vertex
u
    let go :: [Vertex] -> Vertex -> Vertex -> ST s [Bounds]
go [] Vertex
_ Vertex
u = [Bounds] -> ST s [Bounds]
forall (m :: * -> *) a. Monad m => a -> m a
return [(Vertex
u, Vertex
r), (Vertex
r, Vertex
u)]
        go (Vertex
p:[Vertex]
ps) Vertex
k Vertex
u = do
            Vertex
cp <- STUArray s Vertex Vertex -> Vertex -> ST s Vertex
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Vertex Vertex
c Vertex
p
            STUArray s Vertex Vertex -> Vertex -> Vertex -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Vertex Vertex
c Vertex
p (Vertex -> ST s ()) -> Vertex -> ST s ()
forall a b. (a -> b) -> a -> b
$ Vertex
cp Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
- Vertex
1
            ((Vertex
u, Vertex
p)Bounds -> [Bounds] -> [Bounds]
forall a. a -> [a] -> [a]
:) ([Bounds] -> [Bounds])
-> ([Bounds] -> [Bounds]) -> [Bounds] -> [Bounds]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Vertex
p, Vertex
u)Bounds -> [Bounds] -> [Bounds]
forall a. a -> [a] -> [a]
:) ([Bounds] -> [Bounds]) -> ST s [Bounds] -> ST s [Bounds]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> if Vertex
p Vertex -> Vertex -> Bool
forall a. Ord a => a -> a -> Bool
< Vertex
k Bool -> Bool -> Bool
&& Vertex
cp Vertex -> Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== Vertex
1
                                        then [Vertex] -> Vertex -> Vertex -> ST s [Bounds]
go [Vertex]
ps Vertex
k Vertex
p
                                        else STUArray s Vertex Vertex -> Vertex -> ST s Vertex
forall s. STUArray s Vertex Vertex -> Vertex -> ST s Vertex
nxtLeaf STUArray s Vertex Vertex
c (Vertex
k Vertex -> Vertex -> Vertex
forall a. Num a => a -> a -> a
+ Vertex
1) ST s Vertex -> (Vertex -> ST s [Bounds]) -> ST s [Bounds]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Vertex -> Vertex -> ST s [Bounds]) -> Vertex -> ST s [Bounds]
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join ([Vertex] -> Vertex -> Vertex -> ST s [Bounds]
go [Vertex]
ps)
    STUArray s Vertex Vertex -> Vertex -> ST s Vertex
forall s. STUArray s Vertex Vertex -> Vertex -> ST s Vertex
nxtLeaf STUArray s Vertex Vertex
c Vertex
l ST s Vertex -> (Vertex -> ST s [Bounds]) -> ST s [Bounds]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Vertex -> Vertex -> ST s [Bounds]) -> Vertex -> ST s [Bounds]
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join ([Vertex] -> Vertex -> Vertex -> ST s [Bounds]
go [Vertex]
ps)