n-gramを作ってランダムに文章を生成するマルコフもどき

Markov.scala

package net.aruneko.markov

import scala.annotation.tailrec
import scala.io.Source
import scala.util.Random

object Main {
  def main(args: Array[String]): Unit = {
    val trigrams = Source.fromResource("kokoro.txt")
      .getLines().filter(_ != "")
      .map(makeNGram(3, _))
      .flatten.toArray.sorted.groupBy(identity).values
      .map(grams => (grams.length, grams.head)).toArray

    val rand = new Random()
    val sentence = buildSentence(rand, trigrams)
    println(sentence)
  }

  def makeNGram(n: Int, str: String): List[String] = {
    @tailrec
    def inner(n: Int, str: String, ngram: List[String]): List[String] = {
      if (str.length <= n) {
        str :: ngram
      } else {
        inner(n, str.tail, str.take(n) :: ngram)
      }
    }
    inner(n, str, List()).reverse
  }

  def buildSentence(rand: Random, nGram: Array[(Int, String)]): String = {
    @tailrec
    def inner(rand: Random, nGram: Array[(Int, String)], sentence: String): String = {
      val index = rand.nextInt(nGram.length)
      val (_, next_string) = nGram(index)
      if (next_string.contains("。")) {
        sentence + next_string
      } else {
        inner(rand, nGram, sentence + next_string)
      }
    }
    inner(rand, nGram, "")
  }
}