F#: Recursive collect and filter over N-ary Tree

This is hurting my brain!

I want to recurse over a tree structure and collect all instances that match some filter into one list.

Here's a sample tree structure

type Tree =
| Node of int * Tree list

Here's a test sample tree:

let test = 
    Node((1,
            [Node(2,
                    [Node(3,[]);
                    Node(3,[])]);
            Node(3,[])]))

Collecting and filtering over nodes with and int value of 3 should give you output like this:

[Node(3,[]);Node(3,[]);Node(3,[])]

The following recursive function should do the trick:

// The 'f' parameter is a predicate specifying 
// whether element should be included in the list or not 
let rec collect f (Node(n, children) as node) = 
  // Process recursively all children
  let rest = children |> List.collect (collect f)
  // Add the current element to the front (in case we want to)
  if (f n) then node::rest else rest

// Sample usage
let nodes = collect (fun n -> n%3 = 0) tree

The function List.collect applies the provided function to all elements of the list children - each call returns a list of elements and List.collect concatenates all the returned lists into a single one.

Alternatively you could write (this maay help understanding how the code works):

let rest = 
   children |> List.map (fun n -> collect f n)
            |> List.concat

The same thing can be also written using list comprehensions:

let rec collect f (Node(n, children) as node) = 
  [ for m in children do 
      // add all returned elements to the result
      yield! collect f m
    // add the current node if the predicate returns true
    if (f n) then yield node ]

EDIT: Updated code to return nodes as pointed out by kvb.

BTW: It is generally a good idea to show some code that you tried to write so far. This helps people understand what part you didn't understand and so you'll get more helpful answers (and it is also considered as polite)

A more complex tail recursive solution.

let filterTree (t : Tree) (predicate : int -> bool) =
    let rec filter acc = function
        | (Node(i, []) as n)::tail -> 
            if predicate i then filter (n::acc) tail
            else filter acc tail
        | (Node(i, child) as n)::tail -> 
            if predicate i then filter (n::acc) (tail @ child)
            else filter acc (tail @ child)
        | [] -> acc

    filter [] [t]

Just for showing usage of F# Sequences Expression (maybe not the best approach, Tomas's solution more likely better I think):

type Tree =
  | Node of int * Tree list

let rec filterTree (t : Tree) (predicate : int -> bool) =
  seq {
    match t with
    | Node(i, tl) ->
        if predicate i then yield t
        for tree in tl do
          yield! filterTree tree predicate 
  }

let test = Node (1, [ Node(2, [ Node(3,[]); Node(3,[]) ]); Node(3,[]) ])

printfn "%A" (filterTree test (fun x -> x = 3))

When my brain hurts cuz it's stuck up a tree, I try to say what I want as simply and clearly as I can:

Given a tree of info, list all sub-trees matching a predicate (in this case, info = 3).

One straightforward way to do it is to list all nodes of the tree, then filter on the predicate.

type 'info tree = Node of 'info * 'info tree list

let rec visit = function
    | Node( info, [] )  as node -> [ node ]
    | Node( info, children ) as node -> node :: List.collect visit children

let filter predicate tree = 
    visit tree
    |> List.filter (fun (Node(info,_)) -> predicate info)

Here's the tree filter run against the OP's sample data:

let result = filter (fun info -> info = 3) test

One thing that surprised me is how easily the code works for any 'info type with the appropriate predicate:

let test2 = 
    Node(("One",
            [Node("Two",
                    [Node("Three",[Node("Five",[]);Node("Three",[])]);
                    Node("Three",[])]);
            Node("Three",[])]))

let res2 = filter  (fun info -> info = "Three") test2

Alternatively, if you wanted to list the info rather than the sub-trees, the code is breath-takingly simple:

let rec visit = function
    | Node( info, [] )  -> [ info ]
    | Node( info, children ) -> info :: List.collect visit children

let filter predicate tree = 
    visit tree
    |> List.filter predicate

which supports the same queries but only returns the 'info records, not the tree structure:

let result = filter (fun info -> info = 3) test

> val result : int list = [3; 3; 3; 3]

Tomas's approach looks standard, but doesn't quite match your example. If you actually want the list of matching nodes rather than values, this should work:

let rec filter f (Node(i,l) as t) =
  let rest = List.collect (filter f) l
  if f t then t::rest
  else rest

let filtered = filter (fun (Node(i,_)) -> i=3) test

Here is an over engineered solution but it shows seperation of concerns with Partial Active Patterns. This isn't the best example for partial active patterns but it was a fun exercise nonetheless. Match statements are evaluated in order.

let (|EqualThree|_|) = function
    | Node(i, _) as n when i = 3 -> Some n
    | _ -> None

let (|HasChildren|_|) = function
    | Node(_, []) -> None
    | Node(_, children) as n -> Some children
    | _ -> None

let filterTree3 (t : Tree) (predicate : int -> bool) =
    let rec filter acc = function
        | EqualThree(n)::tail & HasChildren(c)::_ -> filter (n::acc) (tail @ c)
        | EqualThree(n)::tail -> filter (n::acc) tail
        | HasChildren(c)::tail -> filter acc (tail @ c)
        | _::tail -> filter acc tail
        | [] -> acc

    filter [] [t]

来源：https://stackoverflow.com/questions/2815236/f-recursive-collect-and-filter-over-n-ary-tree

标签

n-ary-tree