module TreeShapedOrder where
import Prelude hiding (null)
import Data.List (groupBy, sort)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Tree (Tree(..), Forest)
import qualified Data.Tree as Tree
import Util
newtype TSO a = TSO { forall a. TSO a -> Map a (Int, a)
unTSO :: Map a (Int,a) } deriving (TSO a -> TSO a -> Bool
(TSO a -> TSO a -> Bool) -> (TSO a -> TSO a -> Bool) -> Eq (TSO a)
forall a. Eq a => TSO a -> TSO a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => TSO a -> TSO a -> Bool
== :: TSO a -> TSO a -> Bool
$c/= :: forall a. Eq a => TSO a -> TSO a -> Bool
/= :: TSO a -> TSO a -> Bool
Eq, Eq (TSO a)
Eq (TSO a) =>
(TSO a -> TSO a -> Ordering)
-> (TSO a -> TSO a -> Bool)
-> (TSO a -> TSO a -> Bool)
-> (TSO a -> TSO a -> Bool)
-> (TSO a -> TSO a -> Bool)
-> (TSO a -> TSO a -> TSO a)
-> (TSO a -> TSO a -> TSO a)
-> Ord (TSO a)
TSO a -> TSO a -> Bool
TSO a -> TSO a -> Ordering
TSO a -> TSO a -> TSO a
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a. Ord a => Eq (TSO a)
forall a. Ord a => TSO a -> TSO a -> Bool
forall a. Ord a => TSO a -> TSO a -> Ordering
forall a. Ord a => TSO a -> TSO a -> TSO a
$ccompare :: forall a. Ord a => TSO a -> TSO a -> Ordering
compare :: TSO a -> TSO a -> Ordering
$c< :: forall a. Ord a => TSO a -> TSO a -> Bool
< :: TSO a -> TSO a -> Bool
$c<= :: forall a. Ord a => TSO a -> TSO a -> Bool
<= :: TSO a -> TSO a -> Bool
$c> :: forall a. Ord a => TSO a -> TSO a -> Bool
> :: TSO a -> TSO a -> Bool
$c>= :: forall a. Ord a => TSO a -> TSO a -> Bool
>= :: TSO a -> TSO a -> Bool
$cmax :: forall a. Ord a => TSO a -> TSO a -> TSO a
max :: TSO a -> TSO a -> TSO a
$cmin :: forall a. Ord a => TSO a -> TSO a -> TSO a
min :: TSO a -> TSO a -> TSO a
Ord)
empty :: TSO a
empty :: forall a. TSO a
empty = Map a (Int, a) -> TSO a
forall a. Map a (Int, a) -> TSO a
TSO (Map a (Int, a) -> TSO a) -> Map a (Int, a) -> TSO a
forall a b. (a -> b) -> a -> b
$ Map a (Int, a)
forall k a. Map k a
Map.empty
insert :: (Ord a, Eq a) => a -> (Int, a) -> TSO a -> TSO a
insert :: forall a. (Ord a, Eq a) => a -> (Int, a) -> TSO a -> TSO a
insert a
a (Int, a)
b (TSO Map a (Int, a)
o) = Map a (Int, a) -> TSO a
forall a. Map a (Int, a) -> TSO a
TSO (Map a (Int, a) -> TSO a) -> Map a (Int, a) -> TSO a
forall a b. (a -> b) -> a -> b
$ a -> (Int, a) -> Map a (Int, a) -> Map a (Int, a)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert a
a (Int, a)
b Map a (Int, a)
o
fromList :: (Ord a, Eq a) => [(a,(Int,a))] -> TSO a
fromList :: forall a. (Ord a, Eq a) => [(a, (Int, a))] -> TSO a
fromList [(a, (Int, a))]
l = (TSO a -> (a, (Int, a)) -> TSO a)
-> TSO a -> [(a, (Int, a))] -> TSO a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\ TSO a
o (a
a,(Int, a)
b) -> a -> (Int, a) -> TSO a -> TSO a
forall a. (Ord a, Eq a) => a -> (Int, a) -> TSO a -> TSO a
insert a
a (Int, a)
b TSO a
o) TSO a
forall a. TSO a
empty [(a, (Int, a))]
l
parents :: (Ord a, Eq a) => a -> TSO a -> [(Int,a)]
parents :: forall a. (Ord a, Eq a) => a -> TSO a -> [(Int, a)]
parents a
a (TSO Map a (Int, a)
o) = Maybe (Int, a) -> [(Int, a)]
loop (a -> Map a (Int, a) -> Maybe (Int, a)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup a
a Map a (Int, a)
o) where
loop :: Maybe (Int, a) -> [(Int, a)]
loop Maybe (Int, a)
Nothing = []
loop (Just (Int
n,a
b)) = (Int
n,a
b) (Int, a) -> [(Int, a)] -> [(Int, a)]
forall a. a -> [a] -> [a]
: Maybe (Int, a) -> [(Int, a)]
loop (a -> Map a (Int, a) -> Maybe (Int, a)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup a
b Map a (Int, a)
o)
parent :: (Ord a, Eq a) => a -> TSO a -> Maybe (Int,a)
parent :: forall a. (Ord a, Eq a) => a -> TSO a -> Maybe (Int, a)
parent a
a TSO a
t = [(Int, a)] -> Maybe (Int, a)
forall a. [a] -> Maybe a
headMaybe ([(Int, a)] -> Maybe (Int, a)) -> [(Int, a)] -> Maybe (Int, a)
forall a b. (a -> b) -> a -> b
$ a -> TSO a -> [(Int, a)]
forall a. (Ord a, Eq a) => a -> TSO a -> [(Int, a)]
parents a
a TSO a
t
isAncestor :: (Ord a, Eq a) => a -> a -> TSO a -> Maybe Int
isAncestor :: forall a. (Ord a, Eq a) => a -> a -> TSO a -> Maybe Int
isAncestor a
a a
b TSO a
o = Int -> [(Int, a)] -> Maybe Int
forall {t}. Num t => t -> [(t, a)] -> Maybe t
loop Int
0 ((Int
0,a
a) (Int, a) -> [(Int, a)] -> [(Int, a)]
forall a. a -> [a] -> [a]
: a -> TSO a -> [(Int, a)]
forall a. (Ord a, Eq a) => a -> TSO a -> [(Int, a)]
parents a
a TSO a
o)
where loop :: t -> [(t, a)] -> Maybe t
loop t
_ [] = Maybe t
forall a. Maybe a
Nothing
loop t
acc ((t
n,a
a') : [(t, a)]
ps) | a
a' a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
b = t -> Maybe t
forall a. a -> Maybe a
Just (t
acc t -> t -> t
forall a. Num a => a -> a -> a
+ t
n)
| Bool
otherwise = t -> [(t, a)] -> Maybe t
loop (t
acc t -> t -> t
forall a. Num a => a -> a -> a
+ t
n) [(t, a)]
ps
diff :: (Ord a, Eq a) => a -> a -> TSO a -> Maybe Int
diff :: forall a. (Ord a, Eq a) => a -> a -> TSO a -> Maybe Int
diff a
a a
b TSO a
o = Maybe Int -> (Int -> Maybe Int) -> Maybe Int -> Maybe Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ((Int -> Int) -> Maybe Int -> Maybe Int
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ Int
k -> -Int
k) (Maybe Int -> Maybe Int) -> Maybe Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ a -> a -> TSO a -> Maybe Int
forall a. (Ord a, Eq a) => a -> a -> TSO a -> Maybe Int
isAncestor a
b a
a TSO a
o) Int -> Maybe Int
forall a. a -> Maybe a
Just (Maybe Int -> Maybe Int) -> Maybe Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ a -> a -> TSO a -> Maybe Int
forall a. (Ord a, Eq a) => a -> a -> TSO a -> Maybe Int
isAncestor a
a a
b TSO a
o
invert :: (Ord a, Eq a) => TSO a -> Map a [(Int,a)]
invert :: forall a. (Ord a, Eq a) => TSO a -> Map a [(Int, a)]
invert (TSO Map a (Int, a)
o) = (a -> (Int, a) -> Map a [(Int, a)] -> Map a [(Int, a)])
-> Map a [(Int, a)] -> Map a (Int, a) -> Map a [(Int, a)]
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey a -> (Int, a) -> Map a [(Int, a)] -> Map a [(Int, a)]
forall {b} {a}.
Ord b =>
b -> (a, b) -> Map b [(a, b)] -> Map b [(a, b)]
step Map a [(Int, a)]
forall k a. Map k a
Map.empty Map a (Int, a)
o where
step :: b -> (a, b) -> Map b [(a, b)] -> Map b [(a, b)]
step b
son (a
dist, b
father) Map b [(a, b)]
m = ([(a, b)] -> [(a, b)] -> [(a, b)])
-> b -> [(a, b)] -> Map b [(a, b)] -> Map b [(a, b)]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith [(a, b)] -> [(a, b)] -> [(a, b)]
forall a. [a] -> [a] -> [a]
(++) b
son [] (Map b [(a, b)] -> Map b [(a, b)])
-> Map b [(a, b)] -> Map b [(a, b)]
forall a b. (a -> b) -> a -> b
$
([(a, b)] -> [(a, b)] -> [(a, b)])
-> b -> [(a, b)] -> Map b [(a, b)] -> Map b [(a, b)]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith [(a, b)] -> [(a, b)] -> [(a, b)]
forall a. [a] -> [a] -> [a]
(++) b
father [(a
dist, b
son)] Map b [(a, b)]
m
height :: (Ord a, Eq a) => a -> TSO a -> Maybe Int
height :: forall a. (Ord a, Eq a) => a -> TSO a -> Maybe Int
height a
a TSO a
t = do
let m :: Map a [(Int, a)]
m = TSO a -> Map a [(Int, a)]
forall a. (Ord a, Eq a) => TSO a -> Map a [(Int, a)]
invert TSO a
t
let loop :: a -> Maybe Int
loop a
father = do
sons <- a -> Map a [(Int, a)] -> Maybe [(Int, a)]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup a
father Map a [(Int, a)]
m
return $ if null sons then 0 else
maximum $ map (\ (Int
n,a
son) -> Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
0 (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+) (Maybe Int -> Int) -> Maybe Int -> Int
forall a b. (a -> b) -> a -> b
$ a -> Maybe Int
loop a
son) sons
a -> Maybe Int
loop a
a
increasesHeight :: (Ord a, Eq a) => a -> (Int, a) -> TSO a -> Bool
increasesHeight :: forall a. (Ord a, Eq a) => a -> (Int, a) -> TSO a -> Bool
increasesHeight a
_ (Int
n,a
b) TSO a
t = Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
0 Int -> Int
forall a. a -> a
id (a -> TSO a -> Maybe Int
forall a. (Ord a, Eq a) => a -> TSO a -> Maybe Int
height a
b TSO a
t)
leaves :: (Ord a, Eq a) => TSO a -> [a]
leaves :: forall a. (Ord a, Eq a) => TSO a -> [a]
leaves TSO a
o = ((a, [(Int, a)]) -> a) -> [(a, [(Int, a)])] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a, [(Int, a)]) -> a
forall a b. (a, b) -> a
fst ([(a, [(Int, a)])] -> [a]) -> [(a, [(Int, a)])] -> [a]
forall a b. (a -> b) -> a -> b
$ ((a, [(Int, a)]) -> Bool) -> [(a, [(Int, a)])] -> [(a, [(Int, a)])]
forall a. (a -> Bool) -> [a] -> [a]
filter (\ (a
_parent, [(Int, a)]
sons) -> [(Int, a)] -> Bool
forall a. Null a => a -> Bool
null [(Int, a)]
sons) ([(a, [(Int, a)])] -> [(a, [(Int, a)])])
-> [(a, [(Int, a)])] -> [(a, [(Int, a)])]
forall a b. (a -> b) -> a -> b
$ Map a [(Int, a)] -> [(a, [(Int, a)])]
forall k a. Map k a -> [(k, a)]
Map.toList (TSO a -> Map a [(Int, a)]
forall a. (Ord a, Eq a) => TSO a -> Map a [(Int, a)]
invert TSO a
o)
pathesToForest :: (Ord a, Eq a) => [[(Int,a)]] -> Forest (Int, a)
pathesToForest :: forall a. (Ord a, Eq a) => [[(Int, a)]] -> Forest (Int, a)
pathesToForest [] = []
pathesToForest [[(Int, a)]]
ll =
([[(Int, a)]] -> Tree (Int, a))
-> [[[(Int, a)]]] -> [Tree (Int, a)]
forall a b. (a -> b) -> [a] -> [b]
map (\ [[(Int, a)]]
l -> (Int, a) -> [Tree (Int, a)] -> Tree (Int, a)
forall a. a -> [Tree a] -> Tree a
Node ([(Int, a)] -> (Int, a)
forall a. HasCallStack => [a] -> a
head ([[(Int, a)]] -> [(Int, a)]
forall a. HasCallStack => [a] -> a
head [[(Int, a)]]
l))
([[(Int, a)]] -> [Tree (Int, a)]
forall a. (Ord a, Eq a) => [[(Int, a)]] -> Forest (Int, a)
pathesToForest ([[(Int, a)]] -> [Tree (Int, a)])
-> [[(Int, a)]] -> [Tree (Int, a)]
forall a b. (a -> b) -> a -> b
$ ([(Int, a)] -> Bool) -> [[(Int, a)]] -> [[(Int, a)]]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> ([(Int, a)] -> Bool) -> [(Int, a)] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Int, a)] -> Bool
forall a. Null a => a -> Bool
null) ([[(Int, a)]] -> [[(Int, a)]]) -> [[(Int, a)]] -> [[(Int, a)]]
forall a b. (a -> b) -> a -> b
$ ([(Int, a)] -> [(Int, a)]) -> [[(Int, a)]] -> [[(Int, a)]]
forall a b. (a -> b) -> [a] -> [b]
map [(Int, a)] -> [(Int, a)]
forall a. HasCallStack => [a] -> [a]
tail [[(Int, a)]]
l)) ([[[(Int, a)]]] -> [Tree (Int, a)])
-> [[[(Int, a)]]] -> [Tree (Int, a)]
forall a b. (a -> b) -> a -> b
$
([(Int, a)] -> [(Int, a)] -> Bool)
-> [[(Int, a)]] -> [[[(Int, a)]]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (\ [(Int, a)]
l [(Int, a)]
l' -> [(Int, a)] -> (Int, a)
forall a. HasCallStack => [a] -> a
head [(Int, a)]
l (Int, a) -> (Int, a) -> Bool
forall a. Eq a => a -> a -> Bool
== [(Int, a)] -> (Int, a)
forall a. HasCallStack => [a] -> a
head [(Int, a)]
l') [[(Int, a)]]
ll
toForest :: (Ord a, Eq a) => TSO a -> Forest (Int,a)
toForest :: forall a. (Ord a, Eq a) => TSO a -> Forest (Int, a)
toForest TSO a
o = [[(Int, a)]] -> Forest (Int, a)
forall a. (Ord a, Eq a) => [[(Int, a)]] -> Forest (Int, a)
pathesToForest ([[(Int, a)]] -> Forest (Int, a))
-> [[(Int, a)]] -> Forest (Int, a)
forall a b. (a -> b) -> a -> b
$ [[(Int, a)]] -> [[(Int, a)]]
forall a. Ord a => [a] -> [a]
sort ([[(Int, a)]] -> [[(Int, a)]]) -> [[(Int, a)]] -> [[(Int, a)]]
forall a b. (a -> b) -> a -> b
$ (a -> [(Int, a)]) -> [a] -> [[(Int, a)]]
forall a b. (a -> b) -> [a] -> [b]
map (\ a
a -> [(Int, a)] -> [(Int, a)]
forall a. [a] -> [a]
reverse ((Int
0,a
a) (Int, a) -> [(Int, a)] -> [(Int, a)]
forall a. a -> [a] -> [a]
: a -> TSO a -> [(Int, a)]
forall a. (Ord a, Eq a) => a -> TSO a -> [(Int, a)]
parents a
a TSO a
o)) ([a] -> [[(Int, a)]]) -> [a] -> [[(Int, a)]]
forall a b. (a -> b) -> a -> b
$ TSO a -> [a]
forall a. (Ord a, Eq a) => TSO a -> [a]
leaves TSO a
o
instance (Ord a, Eq a, Show a) => Show (TSO a) where
show :: TSO a -> String
show TSO a
o = [Tree String] -> String
Tree.drawForest ([Tree String] -> String) -> [Tree String] -> String
forall a b. (a -> b) -> a -> b
$ (Tree (Int, a) -> Tree String) -> [Tree (Int, a)] -> [Tree String]
forall a b. (a -> b) -> [a] -> [b]
map (((Int, a) -> String) -> Tree (Int, a) -> Tree String
forall a b. (a -> b) -> Tree a -> Tree b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int, a) -> String
forall a. Show a => a -> String
show) ([Tree (Int, a)] -> [Tree String])
-> [Tree (Int, a)] -> [Tree String]
forall a b. (a -> b) -> a -> b
$ TSO a -> [Tree (Int, a)]
forall a. (Ord a, Eq a) => TSO a -> Forest (Int, a)
toForest TSO a
o
l1 :: [(String, (Int, String))]
l1 :: [(String, (Int, String))]
l1 = ((Integer, Integer) -> (String, (Int, String)))
-> [(Integer, Integer)] -> [(String, (Int, String))]
forall a b. (a -> b) -> [a] -> [b]
map (\ (Integer
k,Integer
l) -> (String
"i" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
k, (Int
1, String
"i" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
l))) [(Integer
0,Integer
1),(Integer
1,Integer
2),(Integer
2,Integer
3),(Integer
3,Integer
4)]
[(String, (Int, String))]
-> [(String, (Int, String))] -> [(String, (Int, String))]
forall a. [a] -> [a] -> [a]
++ [(String
"j2",(Int
1,String
"i3"))]
o1 :: TSO String
o1 :: TSO String
o1 = [(String, (Int, String))] -> TSO String
forall a. (Ord a, Eq a) => [(a, (Int, a))] -> TSO a
fromList [(String, (Int, String))]
l1
t1, t2, t3, t4, t5 :: Maybe Int
t1 :: Maybe Int
t1 = String -> String -> TSO String -> Maybe Int
forall a. (Ord a, Eq a) => a -> a -> TSO a -> Maybe Int
diff String
"i2" String
"i1" TSO String
o1
t2 :: Maybe Int
t2 = String -> String -> TSO String -> Maybe Int
forall a. (Ord a, Eq a) => a -> a -> TSO a -> Maybe Int
diff String
"i2" String
"j2" TSO String
o1
t3 :: Maybe Int
t3 = String -> TSO String -> Maybe Int
forall a. (Ord a, Eq a) => a -> TSO a -> Maybe Int
height String
"i2" TSO String
o1
t4 :: Maybe Int
t4 = String -> TSO String -> Maybe Int
forall a. (Ord a, Eq a) => a -> TSO a -> Maybe Int
height String
"i4" TSO String
o1
t5 :: Maybe Int
t5 = String -> TSO String -> Maybe Int
forall a. (Ord a, Eq a) => a -> TSO a -> Maybe Int
height String
"k" TSO String
o1