In Swift, an efficient function that separates an array into 2 arrays based on a predicate

Question

Note: I'm currently still using Swift 2.2, but open to Swift 3 solutions as well

I'm looking to create a function that operates very closely to filter, except that it keeps the non-matching results as well, and maintains sort order. For instance, say you wanted to filter out numbers divisible by 3 in an array and still keep the list of numbers that aren't divisible by 3.

---

Option 1: Using filter

With filter, you only get the list of numbers divisible by 3 and the original list stays unchanged. You can then filter the original list again with the opposite predicate, but this is an unnecessary second pass. The code looks like this:

let numbers = [1,2,3,4,5,6,7,8,9,10] let divisibleBy3 = numbers.filter { $0 % 3 == 0 } // [3,6,9] let theRest = numbers.filter { $0 % 3 != 0 }      // [1,2,4,5,7,8,10] 

It's true that this is pretty readable, but the fact that it does 2 passes seems inefficient to me, especially if the predicate were more complicated. That's twice as many checks as are actually needed.

Option 2: Using custom separate function in Collection extension

My next attempt was extending Collection and making a function I called separate. This function would take the collection and go through the elements one at a time and add them to either the matching list or the non-matching list. The code looks like this:

extension Collection {     func separate(predicate: (Generator.Element) -> Bool) -> (matching: [Generator.Element], notMatching: [Generator.Element]) {         var groups: ([Generator.Element],[Generator.Element]) = ([],[])         for element in self {             if predicate(element) {                 groups.0.append(element)             } else {                 groups.1.append(element)             }         }         return groups     } } 

I can then use the function like this:

let numbers = [1,2,3,4,5,6,7,8,9,10] let groups = numbers.separate { $0 % 3 == 0 } let matching = groups.matching       // [3,6,9] let notMatching = groups.notMatching // [1,2,4,5,7,8,10] 

This is also pretty clean, but the only thing I don't like is that I'm using a tuple as a return type. Maybe others will disagree, but I'd prefer returning the same type as self for chaining. But technically, you can just grab either .matching or .notMatching, which is the same type as self, and you can chain off of either of them.

Option 3: Using a custom, mutating removeIf function in Array extension

My issue with separate returning a tuple led me to try to make a function that would modify self by removing the matches as it found them and adding them to a new list, and returning the matches list at the end. The returned list is your matches, and the array is trimmed of those values. Order is preserved in both arrays. The code looks like this:

extension Array {     mutating func removeIf(predicate: (Element) -> Bool) -> [Element] {         var removedCount: Int = 0         var removed: [Element] = []          for (index,element) in self.enumerated() {             if predicate(element) {                 removed.append(self.remove(at: index-removedCount))                 removedCount += 1             }         }          return removed     } } 

And it is used like this:

var numbers = [1,2,3,4,5,6,7,8,9,10] let divisibleBy3 = numbers.removeIf { $0 % 3 == 0 } // divisibleBy3: [3,6,9] // numbers:      [1,2,4,5,7,8,10] 

This function had to be implemented in an extension of Array, because the concept of removing an element at a particular index doesn't apply to regular Collections (an Array is defined as public struct Array<Element> : RandomAccessCollection, MutableCollection, and it directly defines the remove(at:) function, rather than getting it from inheritance or a protocol).

Option 4: Combination of Option 2 and 3

I'm a big fan of code reuse, and after coming up with Option 3, I realized I could probably reuse the separate function from Option 2. I came up with this:

extension Array {     mutating func removeIf(predicate: (Element) -> Bool) -> [Element] {         let groups = self.separate(predicate: predicate)         self = groups.notMatching         return groups.matching     } } 

And it's used just like in Option 3.

I was concerned with performance, so I ran each Option through XCTest's measure with 1000 iterations. These were the results:

Option 1: 9 ms Option 2: 7 ms Option 3: 10 ms Option 4: 8 ms 

---

Option 5: Based on negaipro's answer

I knew about partition, but I wasn't going to consider it because it didn't preserve the sort order. negaipro's answer is essentially partition, but it got me thinking. The idea with partition is to swap elements that match with the pivot point, thus ensuring that everything on one side of the end pivot point will all match the predicate and the other side doesn't. I took that idea and changed the action to "move to the end". So matches are removed from their spot and added to the end.

extension Array {     mutating func swapIfModified(predicate: (Element) -> Bool) -> Int {         var matchCount = 0         var index = 0         while index < (count-matchCount) {             if predicate(self[index]) {                 append(remove(at: index))                 matchCount += 1             } else {                 index += 1             }         }          return count-matchCount     } } 

Under my initial tests using an array with 10 numbers, it was comparable to the other Options. But I was concerned with the performance of the line append(remove(at: index)). So I tried all the Options again with arrays going from 1 to 1000, and this option was definitely the slowest.

---

Conclusion:

There isn't a big performance difference between these options. And since Option 4 was faster than Option 3 and reuses the code from Option 2, I'm inclined to throw out Option 3. So I'm leaning towards using plain old filter when I don't care about the unfiltered results (and, likewise, for when I don't care about the filtered results, since it's just using the opposite predicate), and then using either separate or removeIf when I care about keeping both the filtered and unfiltered results.

Question:

So, am I missing something built into Swift that does this already? Is there a better way to accomplish this? Is my extension syntax missing anything (anything that could make it apply this concept to more areas, for example)?

Answer 1

let objects: [Int] = Array(1..<11) let split = objects.reduce(([Int](), [Int]())) { (value, object) -> ([Int], [Int]) in      var value = value      if object % 2 == 0 {         value.1.append(object)     } else {         value.0.append(object)     }      return value } 

Answer 2

Swift 4 Solution

partition(by:)

It reorders the origin array and returns start index of subarray satisfies the predicate.

In this example it returns 7.

0..<7 elemets aren't divisible by 3 and 7..n-1 elements are divisible by 3.

 var numbers = [1,2,3,4,5,6,7,8,9,10]  let partition = numbers.partition(by: { $0 % 3 == 0 })  let divisibleBy3 = Array(numbers[..<partition]) //[3,6,9]  let theRest = Array(numbers[partition...]) //[1,2,4,5,7,8,10]