Extension for grouping array into two dimensional array based on element ([] -> [[]])

Question

I'm new at generics, but I've written an array extension to group an array by array element into a two dimensional array:

extension Array {
    func group<U: Hashable>(by key: (Element) -> U) -> [[Element]] {
        //keeps track of what the integer index is per group item
        var indexKeys = [U : Int]()

        var grouped = [[Element]]()
        for element in self {
            let key = key(element)

            if let ind = indexKeys[key] {
                grouped[ind].append(element)
            }
            else {
                grouped.append([element])
                indexKeys[key] = grouped.count - 1
            }
        }
        return grouped
    }
}

For an array of this struct:

struct Thing {
   var category: String
   var name: String
}

I want to be able to group an array of things [Thing] into a two dimensional array [[Thing]] based on Thing.category.

I use the extension like this:

let things = [
   Thing(category: "A", name: "Apple"),
   Thing(category: "B", name: "Boy"),
   Thing(category: "A", name: "Alligator"),
   Thing(category: "B", name: "Ball"),
   Thing(category: "B", name: "Billboard")
]
let groupedThings = things.group { $0.category }

This returns something like:

//groupedThings
[
   [
      Thing(category: "A", name: "Apple"),
      Thing(category: "A", name: "Alligator")
   ],
   [
      Thing(category: "B", name: "Boy"),
      Thing(category: "B", name: "Ball"),
      Thing(category: "B", name: "Billboard")
   ]
]

Please note that I'm not concerned about sort order.

How does my extension look? Since I'm fairly new at this, I'd like to know if I'm over-complicating it. Is there a way to write the extension more concisely? What about speed issues?

Answer 1

That looks like a clean implementation to me. One minor thing: I find it more natural to assign the new index before the group array is extended, so that you don't have to subtract one:

indexKeys[key] = grouped.count
grouped.append([element])

The Swift standard library defines many methods for the Sequence protocol and not for concrete sequences like Array, examples are map, filter, min/max, contains ...

You can do the same with your method to make it more universally applicable:

extension Sequence {
    // ... no other changes necessary ...
}

However: You are essentially reinventing the wheel. In Swift 4 there is a Dictionary(grouping:by) method which does almost what your code does, only that a dictionary is returned:

let groupedThings = Dictionary(grouping: things, by: { $0.category } )
// [String : [Thing]]

The keys are the categories, and the values are arrays of corresponding things. With

Array(groupedThings.values)

you get a nested array which (apart from the order) is identical to what your method returns.