Объединить отсортированный массив в Swift

Question

Я пытался реализовать алгоритм сортировки массивов слияния в Swift для соревнования и столкнулся с проблемой производительности. Мое решение основано на использовании MinHeap и работает хорошо, но требует больше времени, чем требуется для конкурса (более 1 секунды). К сожалению, я не знаю фактический размер тестового ввода. Ограничения следующие: k <= 1024, количество элементов в каждом массиве должно быть меньше k * 10. </p>

Что я делаю не так? Что я могу сделать, чтобы улучшить производительность моего решения?

Мое решение:

// Reading input file
let k = Int(readLine()!)!
var input: [[UInt8]] = Array(repeating: [], count: k)
for i in 0..<k {
    input[i] = readLine()!.split(separator: " ").dropFirst().map{ UInt8($0)! }
}

// Create min heap
var heap: Heap<(indexOfArray: Int, indexWithinArray: Int)> = Heap { [input] (left, right) -> Bool in
    return input[left.indexOfArray][left.indexWithinArray] < input[right.indexOfArray][right.indexWithinArray]
}

// Fill heap with indices of min element in each array
for indexOfArray in 0..<k {
    if input[indexOfArray].isEmpty { continue }
    heap.insert((indexOfArray: indexOfArray, indexWithinArray: 0))
}

while let min = heap.peek() {
    // Print min element
    print(input[min.indexOfArray][min.indexWithinArray])

    // If it's the last element of the array, remove it from the heap
    if min.indexWithinArray + 1 == input[min.indexOfArray].endIndex {
        heap.remove()
    } else {
        // Increment the index and decrease its' priority in the heap if needed
        heap.updateTop((indexOfArray: min.indexOfArray, indexWithinArray: min.indexWithinArray + 1))
    }
}

Оттуда была взята реализация кучи: https://github.com/raywenderlich/swift-algorithm-club/tree/master/Heap

Единственное, что я добавил, - это функция обновления верхнего элемента:

extension Heap {
    mutating func updateTop(_ newElement: T) {
        if isEmpty { return }
        nodes[0] = newElement
        shiftDown(0)
    }
}

Пример ввода:

4
6 2 26 64 88 96 96
4 8 20 65 86
7 1 4 16 42 58 61 69
1 84

// Format note:
// The first line - k - is the number of arrays
// The first element of each line contains the number of elements in k-th array