Elegant way to split an array in swift

Just for fun a generic implementation that would work with strings as well:

extension Collection {
    func every(n: Int, start: Int = 0) -> UnfoldSequence<Element,Index> {
        sequence(state: dropFirst(start).startIndex) { index in
            guard index < endIndex else { return nil }
            defer { index = self.index(index, offsetBy: n, limitedBy: endIndex) ?? endIndex }
            return self[index]
        }
    }
}

---

extension RangeReplaceableCollection {
    func splitIn(subSequences n: Int) -> [SubSequence] {
        (0..<n).map { .init(every(n: n, start: $0)) }
    }
}

---

[0, 1, 2, 3, 4, 5, 6].splitIn(subSequences: 3)   // [[0, 3, 6], [1, 4], [2, 5]]
[0, 1, 2].splitIn(subSequences: 4)               // [[0], [1], [2], []]
"0123456".splitIn(subSequences: 3)               // ["036", "14", "25"]

You can replace both loops with a map() operation:

extension Array {
    func splitInSubArrays(into size: Int) -> [[Element]] {
        return (0..<size).map {
            stride(from: $0, to: count, by: size).map { self[$0] }
        }
    }
}

The outer map() maps each offset to the corresponding array, and the inner map() maps the indices to the array elements.

Examples:

print([0, 1, 2, 3, 4, 5, 6].splitInSubArrays(into: 3))
// [[0, 3, 6], [1, 4], [2, 5]]

print([0, 1, 2].splitInSubArrays(into: 4))
// [[0], [1], [2], []]