How can I use modern kotlin functional techniques to solve this recursive loop problem?

Issue

I have a kotlin problem, can you come up with a elegant way of solving it?

so effectively I have a list of objects which I want to arrange in order based on whether there is a chain. This is the model object

data class Sequence(
    val id: String,
    val previousSequence: List<String>
  )

So the previousSequence tells you whether or not it is in a chain and contains an id of another sequence (or it can be empty). previousSequence is a list of strings but you can assume that the list will only contain one entry. (Don't ask I'm inheriting somebody else poor design)

So effectively the test data below

  Sequence("2b", listOf("1a")),
  Sequence("1a", emptyList()),
  Sequence("3c", listOf("2b")),
  Sequence("91z", emptyList()),
  Sequence("92z", listOf("91z")),
  Sequence("NO_CHAIN", emptyList())

results in a list of lists as follows:

[
  [
   Sequence(id=1a, previousSequence=[]), 
   Sequence(id=2b, previousSequence=[1a]), 
   Sequence(id=3c, previousSequence=[2b])
  ], 
  
  [
   Sequence(id=91z, previousSequence=[]), 
   Sequence(id=92z, previousSequence=[91z])
  ]
]

The below code works.. using recursion. but was hoping for a more elegant solution. Can you come up with one?

data class Sequence(
    val id: String,
    val previousSequence: List<String>
  )

  @Test
  fun `Test chaining of sequence`() {
    val sequences = listOf(
      Sequence("2b", listOf("1a")),
      Sequence("1a", emptyList()),
      Sequence("3c", listOf("2b")),
      Sequence("91z", emptyList()),
      Sequence("92z", listOf("91z")),
      Sequence("NO_CHAIN", emptyList())
    )

    val sequenceChains = sequences.filter { it.previousSequence.isNotEmpty() }
    val baseOfChainList = sequences.filter { base ->
      base.previousSequence.isEmpty() && sequenceChains.any { it.previousSequence.contains(base.id) }
    }

    val listOfChains: MutableList<MutableList<Sequence>> = mutableListOf()

    baseOfChainList.forEach {
      val individualList: MutableList<Sequence> = mutableListOf()
      individualList.add(it)
      individualList.addAll(recursiveAddSequence(it.id, sequenceChains))
      listOfChains.add(individualList)
    }

    println(listOfChains)
  }


  fun recursiveAddSequence(id: String, sequenceChains: List<Sequence>): List<Sequence> {
    val individualChain: MutableList<Sequence> = mutableListOf()
    sequenceChains.forEach {
      if (it.previousSequence.contains(id)) {
        individualChain.add(it)
        individualChain.addAll(recursiveAddSequence(it.id, sequenceChains))
      }
    }
    return individualChain
  }

Solution

We can solve this with the following code:

val (heads, notHeads) = sequences.partition { it.previousSequence.isEmpty() }
val sequencesByPrevious = notHeads.associateBy { it.previousSequence.first() }
val result = heads.map { head ->
    generateSequence(head) { sequencesByPrevious[it.id] }.toList()
}

First, we prepare a map of items by their previous item. Then for each item that is the head (does not have a previous item), we iterate over next elements one step at a time.

generateSequence() is a tricky part. This function creates a lazy and unbound sequence of items (note it is a different "sequence" than in your example). It acquires subsequent items by applying the provided lambda to the previous item and it does this over and over again until it gets null. This single line is an equivalent of this loop:

val list = mutableListOf<Sequence>()
var curr: Sequence? = head
while (curr != null) {
    list += curr
    curr = sequencesByPrevious[curr.id]
}

If we are not interested in elements that don't belong to any chain (they don't reference and are not referenced by any other item), then we can simply filter them out (thanks @Joffrey):

result.filter { it.size > 1 }

As a bonus: if I'm correct, your solution has at least quadratic time complexity. Above solution is linear.

Answered By - broot