Word Count Code Examples

WordCount.scala

/**
  * @author Kevin Lee
  * @since 2016-01-15
  */
object WordCount extends App {
  def wordCount(lines: String): Map[String, Array[Int]] =
    lines.split('\n')
      .map(_.trim)
      .map(_.split("[\\s]+"))
      .zipWithIndex
      .flatMap {
        case (words, index) => words.map((_, index + 1))
      }
      .groupBy {
        case (word, _) => word
      }
      .mapValues(
        _.map {
          case (word, lineNumber) => lineNumber
        } distinct
      )

  def wordCount2(lines: String): Map[String, Array[Int]] =
    (for {
      (line, index) <- lines.split('\n').zipWithIndex
      words = line.trim.split("[\\s]+")
      word <- words
      lineNumber = index + 1
    } yield (word, lineNumber))
    .groupBy {
      case (word, _) => word
    }
    .mapValues(
      _.map {
        case (_, lineNumber) => lineNumber
      } distinct
    )


  case class WordLine(word: String, lineNumber: Int)

  def wordCount3(lines: String): Map[String, Array[Int]] = {
    val wordsAndLines = for {
      (line, index) <- lines.split('\n').zipWithIndex
      words = line.trim.split("[\\s]+")
      word <- words
      lineNumber = index + 1
    } yield WordLine(word, lineNumber)

    wordsAndLines.groupBy(_.word)
                 .mapValues(_.map(_.lineNumber).distinct)
  }

  val foo =
  """World of Warcraft
    |Hello World
    |Hello Hello""".stripMargin

  List(wordCount(foo), wordCount2(foo), wordCount3(foo))
      .foreach { result =>
        result.foreach {
          case (word, lineNumbers) =>
            println(s"$word : ${lineNumbers.mkString(" ")}")
        }
        println("-------------------------------------")
      }
}