-- Copyright 2013 Evan Laforge -- This program is distributed under the terms of the GNU General Public -- License 3.0, see COPYING or http://www.gnu.org/licenses/gpl-3.0.txt -- | Some more utilities for "Data.Tree". module Util.Trees ( edges , paths, flatPaths , find , findAll , findWithParents , leaves , filter ) where import Prelude hiding (filter) import qualified Data.Foldable as Foldable import qualified Data.List as List import Data.Maybe (mapMaybe) import qualified Data.Tree as Tree import Data.Tree (Tree(..)) import Control.Monad -- | The edges of a forest, as (parent, child). edges :: [Tree a] -> [(a, a)] edges :: forall a. [Tree a] -> [(a, a)] edges = (Tree a -> [(a, a)]) -> [Tree a] -> [(a, a)] forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b] concatMap ((Tree a -> [(a, a)]) -> [Tree a] -> [(a, a)]) -> (Tree a -> [(a, a)]) -> [Tree a] -> [(a, a)] forall a b. (a -> b) -> a -> b $ \(Node a val [Tree a] subs) -> [(a val, a sub_val) | Node a sub_val [Tree a] _ <- [Tree a] subs] [(a, a)] -> [(a, a)] -> [(a, a)] forall a. [a] -> [a] -> [a] ++ [Tree a] -> [(a, a)] forall a. [Tree a] -> [(a, a)] edges [Tree a] subs -- | Get every element along with its parents. The parents are closest first -- root last. paths :: [Tree a] -> [(Tree a, [Tree a])] paths :: forall a. [Tree a] -> [(Tree a, [Tree a])] paths = (Tree a -> [(Tree a, [Tree a])]) -> [Tree a] -> [(Tree a, [Tree a])] forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b] concatMap ([Tree a] -> Tree a -> [(Tree a, [Tree a])] forall {a}. [Tree a] -> Tree a -> [(Tree a, [Tree a])] go []) where go :: [Tree a] -> Tree a -> [(Tree a, [Tree a])] go [Tree a] parents tree :: Tree a tree@(Node a _ [Tree a] subs) = (Tree a tree, [Tree a] parents) (Tree a, [Tree a]) -> [(Tree a, [Tree a])] -> [(Tree a, [Tree a])] forall a. a -> [a] -> [a] : (Tree a -> [(Tree a, [Tree a])]) -> [Tree a] -> [(Tree a, [Tree a])] forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b] concatMap ([Tree a] -> Tree a -> [(Tree a, [Tree a])] go (Tree a treeTree a -> [Tree a] -> [Tree a] forall a. a -> [a] -> [a] :[Tree a] parents)) [Tree a] subs -- | Like 'paths' but the parents and children have been flattened. flatPaths :: [Tree a] -> [(a, [a], [a])] -- ^ (element, parents, children). Parents are closest first root last, -- children are in depth first order. flatPaths :: forall a. [Tree a] -> [(a, [a], [a])] flatPaths = ((Tree a, [Tree a]) -> (a, [a], [a])) -> [(Tree a, [Tree a])] -> [(a, [a], [a])] forall a b. (a -> b) -> [a] -> [b] map (Tree a, [Tree a]) -> (a, [a], [a]) forall {b} {b}. (Tree b, [Tree b]) -> (b, [b], [b]) flatten ([(Tree a, [Tree a])] -> [(a, [a], [a])]) -> ([Tree a] -> [(Tree a, [Tree a])]) -> [Tree a] -> [(a, [a], [a])] forall b c a. (b -> c) -> (a -> b) -> a -> c . [Tree a] -> [(Tree a, [Tree a])] forall a. [Tree a] -> [(Tree a, [Tree a])] paths where flatten :: (Tree b, [Tree b]) -> (b, [b], [b]) flatten (Node b val [Tree b] subs, [Tree b] parents) = (b val, (Tree b -> b) -> [Tree b] -> [b] forall a b. (a -> b) -> [a] -> [b] map Tree b -> b forall a. Tree a -> a Tree.rootLabel [Tree b] parents, (Tree b -> [b]) -> [Tree b] -> [b] forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b] concatMap Tree b -> [b] forall a. Tree a -> [a] Tree.flatten [Tree b] subs) -- | Given a predicate, return the first depthwise matching element and -- its children. find :: (a -> Bool) -> [Tree a] -> Maybe (Tree a) find :: forall a. (a -> Bool) -> [Tree a] -> Maybe (Tree a) find a -> Bool p [Tree a] trees = case (Tree a -> Bool) -> [Tree a] -> Maybe (Tree a) forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a List.find (a -> Bool p (a -> Bool) -> (Tree a -> a) -> Tree a -> Bool forall b c a. (b -> c) -> (a -> b) -> a -> c . Tree a -> a forall a. Tree a -> a rootLabel) [Tree a] trees of Maybe (Tree a) Nothing -> [Maybe (Tree a)] -> Maybe (Tree a) forall (t :: * -> *) (m :: * -> *) a. (Foldable t, MonadPlus m) => t (m a) -> m a msum ([Maybe (Tree a)] -> Maybe (Tree a)) -> [Maybe (Tree a)] -> Maybe (Tree a) forall a b. (a -> b) -> a -> b $ (Tree a -> Maybe (Tree a)) -> [Tree a] -> [Maybe (Tree a)] forall a b. (a -> b) -> [a] -> [b] map ((a -> Bool) -> [Tree a] -> Maybe (Tree a) forall a. (a -> Bool) -> [Tree a] -> Maybe (Tree a) find a -> Bool p ([Tree a] -> Maybe (Tree a)) -> (Tree a -> [Tree a]) -> Tree a -> Maybe (Tree a) forall b c a. (b -> c) -> (a -> b) -> a -> c . Tree a -> [Tree a] forall a. Tree a -> [Tree a] subForest) [Tree a] trees Just Tree a tree -> Tree a -> Maybe (Tree a) forall a. a -> Maybe a Just Tree a tree -- | Return the unmodified subtrees that match the predicate. This is like -- 'find', except that it returns all matching subtrees instead of the first -- one. findAll :: (a -> Bool) -> [Tree a] -> [Tree a] findAll :: forall a. (a -> Bool) -> [Tree a] -> [Tree a] findAll a -> Bool f = (Tree a -> [Tree a]) -> [Tree a] -> [Tree a] forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b] concatMap Tree a -> [Tree a] node where node :: Tree a -> [Tree a] node (Node a a [Tree a] trees) | a -> Bool f a a = [a -> [Tree a] -> Tree a forall a. a -> [Tree a] -> Tree a Node a a [Tree a] trees] | Bool otherwise = (Tree a -> [Tree a]) -> [Tree a] -> [Tree a] forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b] concatMap Tree a -> [Tree a] node [Tree a] trees -- | Like 'paths', but return only the element that matches the predicate. findWithParents :: (a -> Bool) -> [Tree a] -> Maybe (Tree a, [Tree a]) findWithParents :: forall a. (a -> Bool) -> [Tree a] -> Maybe (Tree a, [Tree a]) findWithParents a -> Bool f = ((Tree a, [Tree a]) -> Bool) -> [(Tree a, [Tree a])] -> Maybe (Tree a, [Tree a]) forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a List.find (a -> Bool f (a -> Bool) -> ((Tree a, [Tree a]) -> a) -> (Tree a, [Tree a]) -> Bool forall b c a. (b -> c) -> (a -> b) -> a -> c . Tree a -> a forall a. Tree a -> a Tree.rootLabel (Tree a -> a) -> ((Tree a, [Tree a]) -> Tree a) -> (Tree a, [Tree a]) -> a forall b c a. (b -> c) -> (a -> b) -> a -> c . (Tree a, [Tree a]) -> Tree a forall a b. (a, b) -> a fst) ([(Tree a, [Tree a])] -> Maybe (Tree a, [Tree a])) -> ([Tree a] -> [(Tree a, [Tree a])]) -> [Tree a] -> Maybe (Tree a, [Tree a]) forall b c a. (b -> c) -> (a -> b) -> a -> c . [Tree a] -> [(Tree a, [Tree a])] forall a. [Tree a] -> [(Tree a, [Tree a])] paths -- | Get all leaves. The list will never be null. leaves :: Tree a -> [a] leaves :: forall a. Tree a -> [a] leaves = Tree a -> [a] forall (t :: * -> *) a. Foldable t => t a -> [a] Foldable.toList -- | Filter out nodes that don't match. filter :: (a -> Bool) -> Tree a -> Maybe (Tree a) filter :: forall a. (a -> Bool) -> Tree a -> Maybe (Tree a) filter a -> Bool f = Tree a -> Maybe (Tree a) node where node :: Tree a -> Maybe (Tree a) node (Node a a [Tree a] trees) | a -> Bool f a a = Tree a -> Maybe (Tree a) forall a. a -> Maybe a Just (Tree a -> Maybe (Tree a)) -> Tree a -> Maybe (Tree a) forall a b. (a -> b) -> a -> b $ a -> [Tree a] -> Tree a forall a. a -> [Tree a] -> Tree a Node a a ((Tree a -> Maybe (Tree a)) -> [Tree a] -> [Tree a] forall a b. (a -> Maybe b) -> [a] -> [b] mapMaybe Tree a -> Maybe (Tree a) node [Tree a] trees) | Bool otherwise = Maybe (Tree a) forall a. Maybe a Nothing