How to sort an array based on a predefined element order?

Question

I would like to sort the elements of an array based on a custom array that defines the element order. For example, assume the following array:

let arr = ["second", "first", "second", "fourth", "third", "second"]

I tried to create an Array extension to be able to sort this array by:

let sortedArr = arr.sortBy(orderArray: ["first","second","third","fourth"])
// desired output: ["first", "second", "second", "second", "third", "fourth": 

However, the extension does not work correctly:

extension Array {
    public func sortBy<T: Comparable>(orderArray: [T]) -> [T]? {
        guard self.count > 0,
            self.first! as? T != nil else {
                return nil
        }

        let indices = self.map {orderArray.index(of: $0 as! T)! }

        return self.sorted { indices[$0] > indices[$1] } // This doesn’t work
    }
}

Any ideas?

Answer 1

One problem with your code is that self.sorted expects a closure comparing array elements, not indices.

Here is a possible solution which also avoids unnecessary type casts and unwrappings (explanations inline):

extension Array where Element: Equatable {
    public func sortBy(orderArray: [Element]) -> [Element]? {

        // Index of each element in `orderArray`:
        let targetIndices = self.flatMap { orderArray.index(of: $0) }
        // Verify that each array element occurs in `orderArray`:
        guard targetIndices.count == self.count else {
            return nil
        }
        // Sort array indices according to their index in `orderArray`:
        let sortedIndices = self.indices.sorted { targetIndices[$0] < targetIndices[$1] }
        // Rearrange elements accordingly:
        return sortedIndices.map { self[$0] }
    }
}

Example:

let arr = ["second", "first", "second", "fourth", "third", "second"]

if let sortedArr = arr.sortBy(orderArray: ["first","second","third","fourth"]) {
    print(sortedArr)
    // ["first", "second", "second", "second", "third", "fourth"]
}

---

To move array elements which are not contained in orderArray to the end of the sorted result (but preserve their relative order), modify the code slightly to

extension Array where Element: Equatable {
    public func sortBy(orderArray: [Element]) -> [Element] {

        // Index of each element in `orderArray`:
        let targetIndices = self.enumerated().map {
            orderArray.index(of: $1) ?? orderArray.count + $0
        }
        // Sort array indices according to their index in `orderArray`:
        let sortedIndices = self.indices.sorted { targetIndices[$0] < targetIndices[$1] }
        // Rearrange elements accordingly:
        return sortedIndices.map { self[$0] }
    }
}

Example:

let arr = ["x", "second", "first", "y", "second", "fourth", "third", "second", "z"]
let sortedArr = arr.sortBy(orderArray: ["first","second","third","fourth"])
print(sortedArr)
// ["first", "second", "second", "second", "third", "fourth", "x", "y", "z"]

Answer 2

Another alternative: emanate from the orderArray to construct a sorted array

Another alternative (the "other way around" as compared to @MartinR:s answer) is to emanate from orderArray and simply construct the "sorted" array based on the number of occurences of elements in the orderArray in the array to be sorted.

---

Sub-alternative #1: allow an orderArray which is not comprehensive w.r.t. the elements to be sorted

In the example implementation of this alternative below, elements in the array to be sorted that are not present in the orderArray have been left in same relative order as in the original array, but at the end of the part of the array that does get sorted (i.e., after any elements that are included in orderArray). Another alternative could be nil return in case the orderArray does not fully cover all elements in the array to be sorted.

extension Array where Element: Equatable {
    /* help method */
    private func countInstances(of element: Element) -> Int {
        return reduce(0) { $0 + ($1 == element ? 1 : 0) }
    }

    /* sort by ordered array method */
    func sortBy(orderArray: [Element]) -> [Element] {
        // construct sorted array based on elements in orderArray
        return orderArray
            .reduce([]) { $0 + Array(repeating: $1, count: countInstances(of: $1)) } + 
            filter { !orderArray.contains($0) }
    }
}

Example usage:

let arr1 = ["second", "first", "unsortable", "second", 
           "anotherUnsortable", "fourth", "third", "second"]
let arr2 = ["foo", "first", "bar"]
let orderArray = ["first", "second", "third", "fourth"]

print(arr1.sortBy(orderArray: orderArray)) 
    // ["first", "second", "second", "second", "third", "fourth", "unsortable", "anotherUnsortable"]

print(arr2.sortBy(orderArray: orderArray)) 
    // ["first", "foo", "bar"]

---

Sub-alternative #2: for a non-comprehensive orderArray, return nil

In case you'd rather return nil if orderArray is not comprehensive w.r.t. the members of the array to be sorted, yuou can modify the sortBy method above accordingly:

extension Array where Element: Equatable {
    private func countInstances(of element: Element) -> Int {
        return reduce(0) { $0 + ($1 == element ? 1 : 0) }
    }

    func sortBy(orderArray: [Element]) -> [Element]? {
        guard reduce(true, { $0 && orderArray.contains($1) }) else { return nil }
        return orderArray
            .reduce([]) { $0 + Array(repeating: $1, count: countInstances(of: $1)) }
    }
}

Example usage:

let arr1 = ["second", "first", "second", 
            "fourth", "third", "second"]
let arr2 = ["foo", "baz", "bar"]
let orderArray = ["first", "second", "third", "fourth"]

print(arr1.sortBy(orderArray: orderArray)) 
    // Optional(["first", "second", "second", "second", "third", "fourth"])

print(arr2.sortBy(orderArray: orderArray)) 
    // nil