Kotlin: Find Count from Nested set in List (more functional approach)

Question

Below function creates a Map, gets the count of passengers where passengers are > minTrips. The code works completely fine. Please see below

fun List<Trip>.filter(minTrips : Int): Set<Passenger> {
    var passengerMap: HashMap<Passenger, Int> = HashMap()

    this.forEach { it: Trip ->
        it.passengers.forEach { it: Passenger ->
            var count: Int? = passengerMap.get(it)
            if (count == null) {
                count = 1
                passengerMap.put(it, count)
            } else {
                count += 1
                passengerMap.put(it, count)
            }
        }
    }

    val filteredMinTrips: Map<Passenger, Int> = passengerMap.filterValues { it >= minTrips }
    println (" Filter Results = ${filteredMinTrips}")
    return filteredMinTrips.keys
}

Even though this is written in Kotlin, it seems like the code was first written in Java and then converted over to Kotlin. If it was truly written in Kotlin I am sure this wouldnt have been so many lines of code. How can I reduce the lines of Code? What would be a more funtional approach to solve this? What function or functions can I use to extract the Passengers Set directly where Passengers are > minTrips? This is too much of a code and seems crazy. Any pointers would be helpful here.

Answer 1

One way you could do this is to take advantage of Kotlin's flatmap and grouping calls. By creating a list of all passengers on all trips, you can group them, count them, and return the ones that have over a certain number.

Assuming you have data classes like this (essential details only):

data class Passenger(val id: Int)
data class Trip(val passengers: List<Passenger>)

I was able to write this:

fun List<Trip>.frequentPassengers(minTrips: Int): Set<Passenger> =
    this
       .flatMap { it.passengers }
       .groupingBy { it }
       .eachCount()
       .filterValues { it >= minTrips }
       .keys

This is nice because it is a single expression. Going through it, we look at each Trip and extract all of its Passengers. If we had just done map here, we would have List<List<Passenger>>, but we want a List<Passenger> so we flatmap to achieve that. Next, we groupBy the Passenger objects themselves, and call eachCount() on the returned object, giving us a Map<Passenger, Int>. Finally we filter the map down the Passengers we find interesting, and return the set of keys.

Note that I renamed your function, List already has a filter on it, and even though the signatures are different I found it confusing.