How I replaced a nested for loop with flatMap

04 November 2020

This is a short post about how I refactored a nested for loop using flatMap. I had a struct that defined name and quantity properties as shown below:

struct Item {
  let name: String
  let quantity: Int
}

The actual code used a custom type rather than a String, but I’ve used String here to simplify the example.

I also had an array of Items. However, I wanted to create a new array that contained each Item’s name for a count of its quantity. So, given the following:

let myList = [
  Item(name: "Apple", quantity: 2),
  Item(name: "Banana", quantity: 1),
  Item(name: "Carrot", quantity: 3),
  Item(name: "Doughnut", quantity: 1)
]

I wanted a new String array containing the following:

[“Apple”, “Apple”, “Banana”, “Carrot”, “Carrot”, “Carrot”, “Doughnut”]

My first attempt was using for loops:

var nameList: [String] = []
for item in myList {
  for _ in 0..<item.quantity {
    nameList.append(item.name)
  }
}

print(nameList)

However, I did not like the nested for loops. It seemed verbose; I wanted to condense the code. Fortunately, Swift provides two things that make this refactoring possible:

  1. An Array initializer that takes repeating and count parameters: Array(repeating:count:)
  2. The flatMap function

I thought of using map, but map would result in an array of arrays. For example:

let nameListMap = myList.map { Array(repeating: $0.name, count: $0.quantity) }
print(nameListMap)

[[“Apple”, “Apple”], [“Banana”], [“Carrot”, “Carrot”, “Carrot”], [“Doughnut”]]

Not what I wanted. Fortunately, there is flatMap. FlatMap works like the map function, but it flattens an array of arrays. Just what I needed! So, the for loop could be replaced with the following:

let nameList2 = myList.flatMap{Array(repeating: $0.name, count: $0.quantity)}
print(nameList2)

[“Apple”, “Apple”, “Banana”, “Carrot”, “Carrot”, “Carrot”, “Doughnut”]

The above code is doing the following:

  1. For each Item in myList a new Array is created. It is populated with the Item’s name for a count specified by the Item’s quantity.
  2. It then joins the arrays, resulting in a single, flat, array

Hurray!

The kofi logo of a coffee mugLike this? Please buy me a coffee!

Share