Kotlin: Merging pre-sorted sequences into a sorted sequence

Question

In Kotlin 1.4, I have a Sequence of Sequences. The items in the inner Sequences are already sorted by one of their properties (id).

I want to merge them into one Sequence which would also be sorted by the same property.

It's quite trivial algorithm (always taking the smallest of the next items from all sequences). But I was wondering:

Does Kotlin standard library have a stateless sequence merging operation for pre-sorted Sequences? Something like Sequence<Sequence<T>>.flattenSorted(c: Comparator) or so.

Edit:

As some have correctly assumed from the context, I am not looking for flattenSorted(), which is stateful, does not leverage the pre-sort, and for, say, 100 sequences of 1_000_000 elements, it wouldn't perform too well. I've reworded the question.

Answer 1

This is similar to this answer by @Ondra Žižka however this variant is IMO easier to grok and use, because:

1. The type of the comparator passed in to the method is not affected by the implementation details i.e. comparing Map.Entry<Iterator<T>, T> rather than just the relevant value T.

2. It uses the sequence builder to yield values into the returned sequence cleanly.

3. Other minor changes to improve readability.

fun <T> List<Sequence<T>>.mergeSorted(comparator: Comparator<T>): Sequence<T> {
  val iteratorToCurrentValues = map { it.iterator() }
    .filter { it.hasNext() }
    .associateWith { it.next() }
    .toMutableMap()

  val c: Comparator<Map.Entry<Iterator<T>, T>> = Comparator.comparing({ it.value }, comparator)

  return sequence {
    while (iteratorToCurrentValues.isNotEmpty()) {
      val smallestEntry = iteratorToCurrentValues.minWithOrNull(c)!!

      yield(smallestEntry.value)

      if (!smallestEntry.key.hasNext())
        iteratorToCurrentValues.remove(smallestEntry.key)
      else
        iteratorToCurrentValues[smallestEntry.key] = smallestEntry.key.next()
    }
  }
}

Answer 2

If there's nothing such, I've implemented this, for X Sequences, needing just a Map of the size X, and running in O(n).

fun <T> Sequence<Sequence<T>>.mergeSorted(comparator: Comparator<Map.Entry<Iterator<T>, T>>): Sequence<T> {

    // A map of iterators to their next value.
    val nexts = this
        .map { it.iterator() }
        .filter { it.hasNext() }
        .associateBy({it}, { it.next() }
    ).toMutableMap()
   
    return object : Sequence<T> {
        override fun iterator() = object : Iterator<T> {
            override fun hasNext() = nexts.isNotEmpty()

            override fun next(): T {
                val smallest = nexts.minWithOrNull(comparator)
                if (smallest == null)
                    throw NoSuchElementException("No more items. Did you forget to call hasNext()?")

                val toReturn = smallest.value
                if (!smallest.key.hasNext())   // This source is depleted.
                    nexts.remove(smallest.key)
                else
                    nexts[smallest.key] = smallest.key.next()
                
                return toReturn
            }
        }
    }
}