Swift recursive flatmap

gistfile1.swift

func recursiveFlatmap<T, U: AnyObject>(list: [U]) -> [T] {
    var results = [T]()
    results.reserveCapacity(list.count)
    
    for element in list {
        if let subList = element as? [U] {
            results += recursiveFlatmap(subList)
        } else if let element = element as? T {
            results.append(element)
        }
    }
    
    return results
}

let arr: [AnyObject] = [1, [[2]], [3, [4, [5, 6, [7, [8]]]]]]
let flattened: [Int] = recursiveFlatmap(arr) // [1, 2, 3, 4, 5, 6, 7, 8]

// Note that this will throw out any values in the initial multidimensional array that are not of
// what ever type `T` is infered to be. If you want any all elements to be flatmapped regardless of their
// Type, explicitly type `flattened` as `[AnyObject]`, or `[Any]` as best fits your purposes.

// What do you think about this for Swift 2 using Protocol Extensions?
// Would you recommend any other optimizations?

extension CollectionType {
    func recursiveFlatmap<T>() -> [T] {
        var results = [T]()
        for element in self {
            if let sublist = element as? [Self.Generator.Element] {
                results += sublist.recursiveFlatmap()
            } else if let element = element as? T {
                results.append(element)
            }
        }

        return results
    }
}

let x2: [AnyObject] = arr2.recursiveFlatmap()
print(x2)

@StephenDaedalus looks good to me. The only thing I can think of to optimize this would be to reserve storage ahead of time for the results array, since its count should be at least that of the count of the input array (self). results.reserveCapacity(count as! Int) This is a micro-optimization though.

Hey guys, Very useful! Especially when you need to flatten complex xml structures just add .filter and .map and your good to go.

I made a simpler version of this recursiveFlatMap method:

https://gist.github.com/eonist/a21a491496409ed3e203d3f47fdbe35c