GroupBy in Swift 2.0

2018/08/23: A new version of this post is available. Using KeyPath to simplify the call site and reduce possible mistakes.

A while ago when Swift came out I had the idea to play with a function that given an array of objects would return an array of arrays of objects. Basically the idea was transforming the return of an API into an two level nested structure to use in UITableViews with sections.

I even remember asking about it to Chris in an email while reading his book about Functional Programming in Swift. He kindly answered that what I was describing was the function groupBy.

Today I wanted to try to create this simple function using some of the nice features of Swift 2.0.

First of all, having the ability to add this functionality to the CollectionType protocol is very handy, it means that any collection will receive this behaviour.

extension CollectionType {

I’ve defined two typealias just to make things more readable for me:

public typealias ItemType = Self.Generator.Element
public typealias Grouper = (ItemType, ItemType) -> Bool

Here is the groupBy method. It takes a closure that determines if two elements belong to the same group.

public func groupBy(grouper: Grouper) -> [[ItemType]] {
        var result : Array<Array<ItemType>> = []

        var previousItem: ItemType?
        var group = [ItemType]()

        for item in self {

Here I use defer to make sure that the current item is set has the next previousItem at the end of each loop, no matter how loop ends.

defer {previousItem = item}

guard allows us to execute a special branch only the first time. I don’t want the rest of the code to be thinking about this so using guard is a nice way of quickly continuing the loop. It also gives access to the non-Optional previousItem to the rest of the code.

            guard let previous = previousItem else {
                group.append(item)
                continue
            }

Nothing fancy for the rest of the loop.

            if grouper(previous, item) {
                // Item in the same group
                group.append(item)
            } else {
                // New group
                result.append(group)
                group = [ItemType]()
                group.append(item)
            }
        }

        result.append(group)

        return result
    }

If we define some test data:

    struct Person {
        let name: String
        let priority: Int
    }

    let people = [
        Person(name: "Alex", priority: 1),
        Person(name: "Anna", priority: 1),
        Person(name: "Julian", priority: 1),
        Person(name: "Andrea", priority: 2),
        Person(name: "Rob", priority: 2),
        Person(name: "John", priority: 2),
        Person(name: "Javi", priority: 4)
    ]

This works and does what we expect.

    let sectioned = people.groupBy { $0.name.characters.first == $1.name.characters.first }

    [
    [Person(name: "Alex", priority: 1), Person(name: "Anna", priority: 1)],
    [Person(name: "Julian", priority: 1)],
    [Person(name: "Andrea", priority: 2)],
    [Person(name: "Rob", priority: 2)],
    [Person(name: "John", priority: 2),
    Person(name: "Javi", priority: 4)]
    ]"

Groupable

To make this work with other objects we can define a Groupable protocol:

    public protocol Groupable {
        func sameGroupAs(otherPerson: Self) -> Bool
    }

And make the Person type conform to it.

    extension Person: Groupable {

        func sameGroupAs(otherPerson: Person) -> Bool {
            let f = self.name.characters.first
            let s = otherPerson.name.characters.first

            return f == s
        }

    }

    people[0].sameGroupAs(people[1]) // alex & anna -> true
    people[0].sameGroupAs(people[2]) // alex & julian -> false

With this new protocol in place we can use Protocol Extensions to implement a default method that uses this sameGroupAs function.

    extension CollectionType where Self.Generator.Element: Groupable {

        public func group() -> [[Self.Generator.Element]] {
            return self.groupBy { $0.sameGroupAs($1) }
        }

    }

Now every Collection of Groupable objects will have this group() function for free.

Unique groups

If you take a look at the results you will see that the function returns different groups that should be the same, Alex & Anna are in one group but Andre is on another. To solve that we have to sort the input collection before grouping it.

    extension CollectionType where Self.Generator.Element: Comparable {
        public func uniquelyGroupBy(grouper: (Self.Generator.Element, Self.Generator.Element) -> Bool) -> [[Self.Generator.Element]] {
                let sorted = self.sort()
                return sorted.groupBy(grouper)
            }
        }
    }

Now the result is more likely what we want:

    [
    [Person(name: "Alex", priority: 1), Person(name: "Andrea", priority: 2), Person(name: "Anna", priority: 1)],
    [Person(name: "Javi", priority: 4), Person(name: "John", priority: 2), Person(name: "Julian", priority: 1)],
    [Person(name: "Rob", priority: 2)]
    ]

Conclusion

This is just a simple function but it shows all the nice tools that Swift 2.0 give us. Having the ability to create extensions for protocols with default functions and restricting those to specific generics it’s really nice.

In the real world I guess that uniquelyGroupBy should be really the default implementation, maybe with an option to turn the sorting of. But I leave that as an exercise to the reader.

Check the complete Sections Playground.