Safe transpose a list of unequal-length lists in Kotlin.

collections.kt

/**
 * Safe transpose a list of unequal-length lists.
 *
 * Example:
 * transpose(List(List(1, 2, 3), List(4, 5, 6), List(7, 8)))
 * -> List(List(1, 4, 7), List(2, 5, 8), List(3, 6))
 */
fun <E> transpose(xs: List<List<E>>): List<List<E>> {
    // Helpers
    fun <E> List<E>.head(): E = this.first()
    fun <E> List<E>.tail(): List<E> = this.takeLast(this.size - 1)
    fun <E> E.append(xs: List<E>): List<E> = listOf(this).plus(xs)

    xs.filter { it.isNotEmpty() }.let { ys ->
        return when (ys.isNotEmpty()) {
            true -> ys.map { it.head() }.append(transpose(ys.map { it.tail() }))
            else -> emptyList()
        }
    }
}

Hi there, I noticed that this will 'collapse' values if they're missing, for example if the input's 2nd row, 3rd column is missing, the transpose will have an element for the 3rd row, 2nd column.

input:
[1, 2, 3]
[4, 5]
[6, 7, 8]

result:
[1, 4, 6]
[2, 5, 7]
[3, 8]

expected:
[1, 4, 6]
[2, 5, 7]
[3, null, 8]

I wrote an alternative that will replace missing elements with null.

fun <E> transpose2(input: List<List<E>>, missingElementDefault: E? = null): List<List<E?>> {

  val columnIndices = input.indices

  val maxRowSize = input.maxOf { it.size }
  val rowIndices = 0 until maxRowSize

  return columnIndices.map { columnIndex ->
    rowIndices.map { rowIndex ->
      // instead of getting input[column][row], get input[row][column]
      val element = input.getOrNull(rowIndex)?.getOrNull(columnIndex)
      element  ?: missingElementDefault
    }
  }
}

And because I felt like pre-optimising, I used Sequences to defer computation.

fun <E> sequenceTranspose(input: List<List<E>>, missingElementDefault: E? = null): Sequence<Sequence<E?>> =
  sequence {

    val columnIndices = input.indices

    val maxRowSize = input.maxOf { it.size }
    val rowIndices = 0 until maxRowSize

    columnIndices.forEach { columnIndex ->
      val transposedRow = sequence {
        rowIndices.forEach { rowIndex ->
          // instead of getting input[column][row], get input[row][column]
          val element = input.getOrNull(rowIndex)?.getOrNull(columnIndex)
          yield(element ?: missingElementDefault)
        }
      }

      yield(transposedRow)
    }
  }

usage

fun main() {
//  val input = listOf(listOf(1, 2, 3), listOf(4, 5, 6), listOf(7, 8))
  val input = listOf(listOf(1, 2, 3), listOf(4, 5), listOf(6, 7, 8))

  println("\ninput:")
  println(input.joinToString("\n"))

  println("\ntranspose:")
  println(transpose(input).joinToString("\n"))

  println("\ntranspose2:")
  println(transpose2(input).joinToString("\n"))

  println("\nsequenceTranspose:")
  println(
    sequenceTranspose(input).joinToString("\n") { 
      // note: unlike List, Sequence requires a terminal operation, `joinToString()`, to be correctly displayed
      it.joinToString() 
    }
  ) 

  println("---")
}

result

input:
[1, 2, 3]
[4, 5]
[6, 7, 8]

transpose:
[1, 4, 6]
[2, 5, 7]
[3, 8]

transpose2:
[1, 4, 6]
[2, 5, 7]
[3, null, 8]

sequenceTranspose:
1, 4, 6
2, 5, 7
3, null, 8
---