Scala FlatMap дает неверные результаты

Question

Учитывая список документов, я хочу получить пары, которые разделяют хотя бы один токен. Чтобы сделать это, я написал код ниже, который делает это через инвертированный индекс.

object TestFlatMap {
 case class Document(id : Int, tokens : List[String])

 def main(args: Array[String]): Unit = {

   val documents = List(
     Document(1, List("A", "B", "C", "D")),
     Document(2, List("A", "B", "E", "F", "G")),
     Document(3, List("E", "G", "H")),
     Document(4, List("A", "L", "M", "N"))
   )

   val expectedTokensIds = List(("A",1), ("A",2), ("A",4), ("B",1), ("B",2), ("C",1), ("D",1), ("E",2), ("E",3), ("F",2), ("G",2), ("G",3), ("H",3), ("L",4), ("M",4), ("N",4)) //Expected tokens - id tuples
   val expectedCouples = Set((1, 2), (1, 4), (2, 3), (2, 4)) //Expected resulting pairs


   /**
     * For each token returns the id of the documents that contains it
     * */
   val tokensIds = documents.flatMap{ document =>
     document.tokens.map{ token =>
       (token, document.id)
     }
   }

   //Check if the tuples are right
   assert(tokensIds.length == expectedTokensIds.length && tokensIds.intersect(expectedTokensIds).length == expectedTokensIds.length, "Error: tokens-ids not matches")

   //Group the documents by the token
   val docIdsByToken = tokensIds.groupBy(_._1).filter(_._2.size > 1)

   /**
     * For each group of documents generate the pairs
     * */
   val couples = docIdsByToken.map{ case (token, docs) =>
     docs.combinations(2).map{ c =>
       val d1 = c.head._2
       val d2 = c.last._2

       if(d1 < d2){
         (d1, d2)
       }
       else{
         (d2, d1)
       }
     }
   }.flatten.toSet


   /**
     * Same operation, but with flatMap
     * For each group of documents generate the pairs
     * */
   val couples1 = docIdsByToken.flatMap{ case (token, docs) =>
     docs.combinations(2).map{ c =>
       val d1 = c.head._2
       val d2 = c.last._2

       if(d1 < d2){
         (d1, d2)
       }
       else{
         (d2, d1)
       }
     }
   }.toSet

   //The results obtained with flatten pass the test
   assert(couples.size == expectedCouples.size && couples.intersect(expectedCouples).size == expectedCouples.size, "Error: couples not matches")
   //The results obtained with flatMap do not pass the test: they are wrong
   assert(couples1.size == expectedCouples.size && couples1.intersect(expectedCouples).size == expectedCouples.size, "Error: couples1 not matches")
}

Проблема в том, что flatMap, который должен генерировать окончательные результаты, не работает должным образом, он возвращает только две пары: (2,3) и (1,2). Я не понимаю, почему это не работает, более того, IntelliJ предлагает мне использовать flatMap вместо использования map, а затем flatten.

Кто-то может объяснить мне, где проблема? Поскольку я не могу понять, у меня также была эта проблема в прошлом.

Спасибо

Лука

Answer 1

Это отличный пример, который демонстрирует, что все хорошие законы монады не обязательно выполняются, если вы переключаетесь между разными типами коллекций во время map/flatMap/flatten,

---

Вы должны конвертировать Map в List, чтобы ключи не перезаписывались повторно при создании другого Map в качестве промежуточного результата, потому что Map переопределит ключи вместо сбора всех пар:

val couples1 = docIdsByToken.toList.flatMap{ case (token, docs) =>
  docs.combinations(2).map{ c =>
    val d1 = c.head._2
    val d2 = c.last._2

    if(d1 < d2){
      (d1, d2)
    }
    else{
      (d2, d1)
    }
  }
}.toSet

Вот гораздо более короткая версия, которая демонстрирует ту же проблему:

val m = Map("A" -> (2, 1), "B" -> (2, 3))
val s = m.flatMap{ case (k, v) => List(v) }.toSet
println(s)

Вместо Set((2, 1), (2, 3))будет производить Set((2, 3))потому что после flatMap и до toSet промежуточный результат - новыйMapи эта карта может содержать только одно значение для key = 2,

Разница с первой версией в том, что после map, вы получите что-то вроде Iterable[List[(Int, Int)]], который не является Mapи, следовательно, не может потерять / переопределить любые ключи.