192 lines
6.4 KiB
Haskell
192 lines
6.4 KiB
Haskell
{-
|
|
Helpers for the manipulation of rose trees.
|
|
Throughout, the _trunk_ refers to successive first children in the tree.
|
|
For example, in the tree
|
|
|
|
> Node a [ Node b [], Node c [Node d []] ]
|
|
|
|
the nodes in the trunk are [a,b].
|
|
-}
|
|
module TreeHelp
|
|
( module Data.Tree
|
|
, module Data.Tree.Lens
|
|
, applyToSubtree
|
|
, applyToSubforest
|
|
, treeFromPost
|
|
, treeFromTrunk
|
|
, splitTrunk
|
|
, applyToRandomNode
|
|
, addToTrunk
|
|
, inorderNumberTree
|
|
, updateSingleNodes
|
|
, updateAllNodes
|
|
, updateRandNode
|
|
, safeUpdateSingleNode
|
|
, msafeUpdateSingleNode
|
|
) where
|
|
import RandomHelp
|
|
|
|
import Data.Maybe
|
|
import Data.Tree.Lens
|
|
import Data.Tree
|
|
import Control.Lens
|
|
|
|
{- | Creates a linear tree.
|
|
Safe. -}
|
|
treeFromPost :: [a] -> a -> Tree a
|
|
treeFromPost xs = treeFromTrunk xs . pure
|
|
|
|
{- | Creates a tree with one trunk branch,
|
|
input as a list, that ends in another tree. -}
|
|
treeFromTrunk
|
|
:: [a] -- ^ The trunk
|
|
-> Tree a -- ^ The end of the tree
|
|
-> Tree a
|
|
treeFromTrunk = flip $ foldr f
|
|
where
|
|
f x t = Node x [t]
|
|
|
|
-- find use for?
|
|
---- | Consider defining this using generalised recursion patterns
|
|
--treeSize :: Tree a -> Int
|
|
--treeSize = length . flatten
|
|
|
|
{- | Applies a function to a specific node determined by a list of indices.
|
|
Unsafe (partial function). -}
|
|
applyToNode :: [Int] -> (a -> a) -> Tree a -> Tree a
|
|
applyToNode is = applyToSubtree is . over root
|
|
|
|
{- | Applies a function to a specific subtree determined by a list of indices.
|
|
Unsafe (partial function). -}
|
|
applyToSubtree :: [Int] -> (Tree a -> Tree a) -> Tree a -> Tree a
|
|
applyToSubtree [] f t = f t
|
|
applyToSubtree (i:is) f (Node x xs) = Node x (xs & ix i %~ applyToSubtree is f)
|
|
|
|
{- | Applies a function to a specific subforest determined by a list of indices.
|
|
Unsafe (partial function). -}
|
|
applyToSubforest :: [Int] -> ([Tree a] -> [Tree a]) -> Tree a -> Tree a
|
|
applyToSubforest is = applyToSubtree is . over branches
|
|
--applyToSubforest [] f (Node p cs) = Node p (f cs)
|
|
--applyToSubforest (i:is) f (Node x xs) = Node x (ys ++ [applyToSubforest is f z] ++ zs)
|
|
-- where
|
|
-- (ys, z:zs) = splitAt i xs
|
|
|
|
-- do not delete: find use for
|
|
--{- |
|
|
--Applies a function to the first node along a trunk that satisfies a given property.
|
|
---}
|
|
--applyToSubTrunkBy :: (a -> Bool) -> (Tree a -> Tree a) -> Tree a -> Tree a
|
|
--applyToSubTrunkBy cond f (Node x (t:ts))
|
|
-- | cond x = f (Node x (t:ts))
|
|
-- | otherwise = Node x (applyToSubTrunkBy cond f t : ts)
|
|
--applyToSubTrunkBy _ _ t = t
|
|
|
|
updateAllNodes :: (a -> Bool) -> (Tree a -> Tree a) -> Tree a -> Tree a
|
|
updateAllNodes f update t@(Node x ts)
|
|
| f x = update t
|
|
| otherwise = updateChildren
|
|
where
|
|
updateChildren = Node x (map (updateAllNodes f update) ts)
|
|
|
|
|
|
-- gives the list of all updates to a single node
|
|
-- there must be a better way of doing something like this
|
|
updateSingleNodes :: (a -> Bool) -> (Tree a -> Tree a) -> Tree a -> [Tree a]
|
|
updateSingleNodes f update t@(Node x ts)
|
|
| f x = update t : updateChildren
|
|
| otherwise = updateChildren
|
|
where
|
|
updateChildren = map (Node x) (subMap (updateSingleNodes f update) ts)
|
|
|
|
mupdateSingleNodes :: Monad m => (a -> Bool) -> (Tree a -> m (Tree a)) -> Tree a -> [m (Tree a)]
|
|
mupdateSingleNodes f update t@(Node x ts)
|
|
| f x = update t : updateChildren
|
|
| otherwise = updateChildren
|
|
where
|
|
updateChildren = map (fmap $ Node x) (msubMap (mupdateSingleNodes f update) ts)
|
|
|
|
updateRandNode :: RandomGen g => (a -> Bool) -> (Tree a -> Tree a) -> Tree a -> State g (Tree a)
|
|
updateRandNode t f = takeOne . updateSingleNodes t f
|
|
|
|
safeUpdateSingleNode :: (a -> Bool) -> (Tree a -> Tree a) -> Tree a -> Tree a
|
|
safeUpdateSingleNode f g t = fromMaybe t $ listToMaybe $ updateSingleNodes f g t
|
|
|
|
msafeUpdateSingleNode :: Monoid b => (a -> Bool) -> (Tree a -> (b, Tree a)) -> Tree a -> (b,Tree a)
|
|
msafeUpdateSingleNode f g t = fromMaybe (mempty,t) $ listToMaybe $ mupdateSingleNodes f g t
|
|
|
|
msubMap :: Functor m => (a -> [m a]) -> [a] -> [m [a]]
|
|
msubMap f (x:xs) = (f x <&> fmap (: xs)) ++ ( (fmap (x :)) <$> msubMap f xs )
|
|
msubMap _ [] = []
|
|
|
|
subMap :: (a -> [a]) -> [a] -> [[a]]
|
|
subMap f (x:xs) = (f x <&> (: xs)) ++ ( (x :) <$> subMap f xs )
|
|
subMap _ [] = []
|
|
|
|
-- find use for?
|
|
--zipTree :: Tree a -> Tree b -> Tree (a,b)
|
|
--zipTree (Node x xs) (Node y ys) = Node (x,y) $ zipWith zipTree xs ys
|
|
|
|
{- | Makes each node into its child number, i.e. the index it has
|
|
in the list of children of its parent. -}
|
|
treeChildNums :: Tree a -> Tree Int
|
|
treeChildNums = setRoot 0
|
|
where
|
|
setRoot :: Int -> Tree a -> Tree Int
|
|
setRoot i (Node _ xs) = Node i (zipWith setRoot [0..] xs)
|
|
|
|
{- | Makes each node into its path, i.e. the list of indices that,
|
|
when followed from the root, lead to the node. -}
|
|
treePaths :: Tree a -> Tree [a]
|
|
treePaths (Node x xs) = (x :) <$> Node [] (map treePaths xs)
|
|
|
|
{- | Picks a random path in the tree.
|
|
Uniform probability that the path leads to any specific node. -}
|
|
randomPath :: RandomGen g => Tree a -> State g [Int]
|
|
randomPath = takeOne . flatten . treePaths . treeChildNums
|
|
|
|
{- | Apply a function to the value of a node;
|
|
the node is picked uniformly at random. -}
|
|
applyToRandomNode :: RandomGen g => (a -> a) -> Tree a -> State g (Tree a)
|
|
applyToRandomNode f t = do
|
|
p <- randomPath t
|
|
return $ applyToNode p f t
|
|
|
|
{- | Add a forest to the end of a tree (along the trunk). -}
|
|
addToTrunk :: Tree a -> [Tree a] -> Tree a
|
|
addToTrunk (Node x []) f = Node x f
|
|
addToTrunk (Node x (t:ts)) f = Node x (addToTrunk t f : ts)
|
|
|
|
{- | Find the depth of a tree along the trunk. -}
|
|
trunkDepth :: Tree a -> Int
|
|
trunkDepth (Node _ []) = 0
|
|
trunkDepth (Node _ (x:_)) = trunkDepth x + 1
|
|
|
|
{- | Split a tree at a given point along its trunk. -}
|
|
splitTrunkAt
|
|
:: Int -- ^ Split depth
|
|
-> Tree a -> (Tree a, [Tree a])
|
|
splitTrunkAt 0 (Node x xs) = (Node x [],xs)
|
|
splitTrunkAt i (Node y (x:xs)) =
|
|
let (t, ts) = splitTrunkAt (i-1) x
|
|
in (Node y (t : xs) , ts)
|
|
splitTrunkAt _ (Node _ []) = error "Trying to split to short a trunk"
|
|
|
|
{- | Split a tree at a random point along its trunk. -}
|
|
splitTrunk :: RandomGen g => Tree a -> State g (Tree a, [Tree a])
|
|
splitTrunk t = do
|
|
i <- state $ randomR (0, trunkDepth t)
|
|
return $ splitTrunkAt i t
|
|
|
|
-- untested
|
|
inorderNumberTree :: Tree a -> Tree (a,Int)
|
|
inorderNumberTree = fst . f 0
|
|
where
|
|
f i (Node x ts) =
|
|
let (ts',i') = g (i+1) ts
|
|
in (Node (x,i) ts', i')
|
|
g i (t:ts) =
|
|
let (t',i') = f i t
|
|
(ts',i'') = g i' ts
|
|
in (t': ts', i'')
|
|
g i [] = ([], i)
|