Return double index of collection's element while iterating

Question

In Kotlin documentation I found the following example:

for ((index, value) in array.withIndex()) {
    println("the element at $index is $value")
}

Is it possible (and how) to do the similar with 2D matrix:

for ((i, j, value) in matrix2D.withIndex()) {
    // but iterate iver double index: i - row, j - column
    if (otherMatrix2D[i, j] > value) doSomething()
}

How to make support this functionality in Kotlin class?

Answer 1

While the solutions proposed by miensol and hotkey are correct it would be the least efficient way to iterate a matrix. For instance, the solution of hotkey makes M * N allocations of Cell<T> plus M allocations of List<Cell<T>> and IntRange plus one allocation of List<List<Cell<T>>> and IntRange. Moreover lists resize when new cells are added so even more allocations happen. That's too much allocations for just iterating a matrix.

Iteration using an inline function

I would recommend you to implement a very similar and very effective at the same time extension function that will be similar to Array<T>.forEachIndexed. This solution doesn't do any allocations at all and as efficient as writing nested for cycles.

inline fun <T> Matrix<T>.forEachIndexed(callback: (Int, Int, T) -> Unit) {
  for (i in 0..cols - 1) {
    for (j in 0..rows - 1) {
      callback(i, j, this[i, j])
    }
  }
}

You can call this function in the following way:

matrix.forEachIndexed { i, j, value ->
  if (otherMatrix[i, j] > value) doSomething()
}

Iteration using a destructive declaration

If you want to use a traditional for-loop with destructive declaration for some reason there exist a way more efficient but hacky solution. It uses a sequence instead of allocating multiple lists and creates only a single instance of Cell, but the Cell itself is mutable.

data class Cell<T>(var i: Int, var j: Int, var value: T)

fun <T> Matrix<T>.withIndex(): Sequence<Cell<T>> {
  val cell = Cell(0, 0, this[0, 0])
  return generateSequence(cell) { cell ->
    cell.j += 1
    if (cell.j >= rows) {
      cell.j = 0
      cell.i += 1
      if (cell.i >= cols) {
        return@generateSequence null
      }
    }
    cell.value = this[cell.i, cell.j]
    cell
  }
}

And you can use this function to iterate a matrix in a for-loop:

for ((i, j, item) in matrix.withIndex()) {
   if (otherMatrix[i, j] > value) doSomething()
}

This solution is lightly less efficient than the first one and not so robust because of a mutable Cell, so I would really recommend you to use the first one.