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,[])]

Solution

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)