Find and return first match in nested lists in Kotlin?

Question

Consider the following two classes:

class ObjectA(val objectBs: List<ObjectB>,
              val otherFields: Any)

class ObjectB(val key: String,
              val otherFields: Any)

The task is to find and return the first ObjectB with a certain key in a List of ObjectA.

Just achieving the goal is simple enough, but doing it nicely and efficiently seems rather tricky. I can't find anything like a "firstIn" or "findIn" function that would allow me to return another type than ObjectA when iterating on a list of ObjectA.

I have a few approaches, one of which looks pretty nice, but is very inefficient:

listOfA.mapNotNull { 
    it.objectBs.firstOrNull { 
        item -> item.key == wantedKey
   } 
}.firstOrNull()

The obvious inefficiency of this code is that it will not stop iterating through listOfA when it has found a match (and there can only be one match, just to be clear). Approaches using filter or find have similar problems, requiring redundant iterations through at least one list of ObjectB.

Is there something in kotlins standard library that would cover such a use case?

Answer 1

If you want an elegant solution you can just do a flatMap like this:

val result: ObjectB? = listOfA.flatMap { it.objectBs }.firstOrNull { it.key == "myKey" }

If you want the efficiency you can do something like this:

val result: ObjectB? = objectAs.firstOrNull {
    it.objectBs.map(ObjectB::key).contains("myKey")
}?.objectBs?.firstOrNull { it.key == "myKey" }

You can also wrap these in an Optional and put it in a function so the users of this operation can have a clean API:

fun List<ObjectA>.findFirstObjectB(key: String): Optional<ObjectB> {
    return Optional.ofNullable(firstOrNull {
        it.objectBs.map(ObjectB::key).contains(key)
    }?.objectBs?.firstOrNull { it.key == key })
}

Answer 2

By converting all the nested elements to a flattened Sequence, they can be iterated lazily, and the overhead of unnecessary iteration is eliminated. This trick is done by combining asSequence and flatMap:

listOfA.asSequence().flatMap { it.objectBs.asSequence() }.find { it.key == wantedKey }

---

I wrote and ran the following code to ensure that it works as expected:

class PrintSequenceDelegate<out T>(private val wrappedSequence: Sequence<T>) : Sequence<T> by wrappedSequence {
    override fun iterator(): Iterator<T> {
        val wrappedIterator = wrappedSequence.iterator()
        return object : Iterator<T> by wrappedIterator {
            override fun next(): T =
                wrappedIterator.next().also { println("Retrieving: $it") }
        }
    }
}

fun <T> Sequence<T>.toPrintDelegate() = PrintSequenceDelegate(this)

fun main() {
    val listOfLists = List(3) { i -> List(3) { j -> "$i$j" } }
    println("List of lists: $listOfLists")
    val found = listOfLists.asSequence().toPrintDelegate().flatMap { it.asSequence().toPrintDelegate() }.find { it == "11" }
    println(if (found != null) "Found: $found" else "Not found")
}

Output:

List of lists: [[00, 01, 02], [10, 11, 12], [20, 21, 22]]
Retrieving: [00, 01, 02]
Retrieving: 00
Retrieving: 01
Retrieving: 02
Retrieving: [10, 11, 12]
Retrieving: 10
Retrieving: 11
Found: 11

Thus we see that the elements (12) after the element found in the containing nested list are not iterated, neither are the following nested lists ([20, 21, 22]).