Выравнивание многоуровневого массива в Swift (Flattening of multi level Array in Swift)

Всем привет! На собеседованиях время от времени встречается вопрос как выравнить многоуровневый массив.

Если это простой двухуровневый массив вида [[1, 2], [3, 4]], то нам поможет стандартная функция высшего порядка flatMap() хоть она и деприкейтнута, но все еще доступна для использования.

Но что если нам нужно нужно выровнять более сложные массивы, например:

[1, [2]] -> [1, 2]

[1, [2, 3]] -> [1, 2, 3]

[1, [2, [3, [4]]]] -> [1, 2, 3, 4]

В этом случае мы можем написать свою рекурсивную функцию:

func flattFunction(array:[Any]) -> [Int] {
    var flattenArray = [Int]()
    
    for element in array {
        if let element = element as? Int {
            flattenArray.append(element)
        } else if let element = element as? [Any] {
            let result = flattFunction(array: element)
            flattenArray += result
        }
    }
    
    return flattenArray
}

Протестируем:

let arr1 = [1]
let arr2 = [1, 2, 3]
let arr3: [Any] = [1, [2]]
let arr4: [Any] = [1, [2, 3]]
let arr5: [Any] = [1, [2, [3, [4]]]]

func testFlattFunc() {
    print(flattFunction(array: arr1))
    print(flattFunction(array: arr2))
    print(flattFunction(array: arr3))
    print(flattFunction(array: arr4))
    print(flattFunction(array: arr5))
}

testFlattFunc()

// Вывод:
// [1]
// [1, 2, 3]
// [1, 2]
// [1, 2, 3]
// [1, 2, 3, 4]

Надеюсь это небольшая заметка окажется полезной.