I've written this code in Scala that extracts bigram statistics from the Reuters dataset. It puts the statistics in a Map[String, SortedMap[String, Int]]. For example, the bigrams for "hi" could look like this:

Map("hello" -> SortedMap("there" -> 10, "friend" -> 6))

which means that "there" follows "hi" 10 times in the dataset, and "friend" follows "hi" 6 times.

My code takes a very long time to run, 22 minutes just for one file of 1.3 megabytes.

The code looks like this:

import java.io.File

import scala.annotation.tailrec
import scala.collection.mutable
import scala.io.Source
import scala.util.matching.Regex
import scala.collection.mutable

object Main extends App {

  val bigrams: Bigrams = Bigrams.fromPath("src/test/resources/mini.sgm")

case class Bigrams(bigrams: BigramsMap) {

  def extractStatistics(path: String): Bigrams = {
    val entry: File = new File(path)
    if (entry.exists && entry.isDirectory) {
      println("Extracting bigrams from " + entry.getPath + "/")
      val bigramsFromDir: BigramsMap = entry
        .filter(file => file.isFile && file.getName.endsWith(".sgm"))
      val bigramsFromSubDirs: Bigrams = entry
        .filter(entry => entry.isDirectory)
        .map(entry => extractStatistics(entry.getAbsolutePath))
    } else if (entry.exists && entry.isFile) {
    } else
    throw new RuntimeException("Incorrect path")

  def getFreqs(word: String): Option[mutable.SortedMap[String, Int]] = {

  def mergeIn(bigramsIn: BigramsMap): Bigrams = {
    Bigrams(Bigrams.merge(bigrams, bigramsIn))

object Bigrams {

  val BODY: Regex = "(?s).*<BODY>(.*)</BODY>(?s).*".r

  def apply(): Bigrams = {
    new Bigrams(BigramsMap())

  def fromPath(path: String): Bigrams = {
    new Bigrams(BigramsMap()).extractStatistics(path)

  // Return a list with the markup for each article
  def readArticles(remainingLines: List[String], acc: List[String]): List[String] = {
    if (remainingLines.size == 1) (acc.head + "\n" + remainingLines.head) +: acc.tail
    else {
      val nextLine = remainingLines.head
      if (nextLine.startsWith("<REUTERS ")) readArticles(remainingLines.tail, nextLine +: acc)
      else readArticles(remainingLines.tail, (acc.head + "\n" + nextLine) +: acc.tail)

  def addBigramsFrom(tokens: List[String], oldBigrams: BigramsMap): BigramsMap = {
    val bigramPairs: List[(String, String)] = Bigrams.getBigrams(tokens)
    val newBigrams: BigramsMap = getBigramsFrom(bigramPairs)
    merge(oldBigrams, newBigrams)

  def merge(bigrams1: Bigrams, bigrams2: Bigrams): Bigrams = {
    Bigrams(merge(bigrams1.bigrams, bigrams2.bigrams))

  def merge(bigrams1: BigramsMap, bigrams2: BigramsMap): BigramsMap = {
    bigrams1.map ++ bigrams2.map
      .map(entry1 => entry1._1 -> (entry1._2 ++ bigrams1.getOrElse(entry1._1, mutable.SortedMap[String, Int]())
      .map(entry2 => entry2._1 -> (entry2._2 + entry1._2.getOrElse(entry2._1, 0))))))

  def getBigramsFrom(path: File): BigramsMap = {
    println("Extracting bigrams from " + path.getPath)
    val file = Source.fromFile(path)
    val fileLines: List[String] = file.getLines().toList
    val articles: List[String] = Bigrams.readArticles(fileLines.tail, List())
    val bodies: List[String] = articles.map(extractBody).filter(body => !body.isEmpty)
    val sentenceTokens: List[List[String]] = bodies.flatMap(getSentenceTokens)
    sentenceTokens.foldLeft(BigramsMap())((acc, tokens) => addBigramsFrom(tokens, acc))

  def getBigramsFrom(tokens: List[(String, String)]): BigramsMap = {

  def getBigrams(tokens: List[String]): List[(String, String)] = {
    map(i => {
        if (i < tokens.size - 1) (tokens(i), tokens(i + 1))
        else null
      .filter(_ != null).toList

  // Return the body of the markup of one article
  def extractBody(article: String): String = {
    try {
      val body: String = article match {
        case Bigrams.BODY(bodyGroup) => bodyGroup
    catch {
      case _: MatchError => ""

  def getSentenceTokens(text: String): List[List[String]] = {
    val separatedBySpace: List[String] = text
      .replace('\n', ' ')
      .replaceAll(" +", " ") // regex
      .split(" ")
      .map(token => if (token.endsWith(",")) token.init.toString else token)

    val splitAt: List[Int] = separatedBySpace.indices
      .filter(i => i > 0 && separatedBySpace(i - 1).endsWith(".") || i == 0)

      .map(sentenceTokens => sentenceTokens.init :+ sentenceTokens.last.substring(0, sentenceTokens.last.length - 1))
      .map(sentenceTokens => sentenceTokens.map(sentenceToken => sentenceToken.toLowerCase))

  def groupBySentenceTokens(tokens: List[String], splitAt: List[Int], sentences: List[List[String]]): List[List[String]] = {
    if (splitAt.size <= 1) {
      if (splitAt.size == 1) {
        sentences :+ tokens.slice(splitAt.head, tokens.size)
      } else {
    else groupBySentenceTokens(tokens, splitAt.tail, sentences :+ tokens.slice(splitAt.head, splitAt.tail.head))

case class BigramsMap(map: Map[String, mutable.SortedMap[String, Int]]) { 
  def addAll(bigrams: List[(String, String)]): BigramsMap = {
    if (bigrams.isEmpty) this
    else {
      val first = bigrams.head._1
      val second = bigrams.head._2
      if (map.contains(first)) {
        if (map(first).contains(second)) {
          BigramsMap(map.updated(first, map(first) + (second -> (map(first)(second) + 1)))).addAll(bigrams.tail)
        } else {
          BigramsMap(map.updated(first, map(first) + (second -> 1))).addAll(bigrams.tail)
      } else {
        BigramsMap(map.updated(first, mutable.SortedMap(second -> 1))).addAll(bigrams.tail)

  def get(key: String): Option[mutable.SortedMap[String, Int]] = map.get(key)

  def getOrElse[V1 >: mutable.SortedMap[String, Int]](key: String, default: V1): V1 = map.getOrElse(key, default)

object BigramsMap {
  def apply(elems: (String, mutable.SortedMap[String, Int])*): BigramsMap = BigramsMap(Map(elems: _*))
  def apply(): BigramsMap = BigramsMap(Map[String, mutable.SortedMap[String, Int]]())

Here is a small test file to run it on. Obviously, the Reuters dataset is much, much bigger.


<!DOCTYPE lewis SYSTEM "lewis.dtd">
<DATE> 9-APR-1987 09:40:15.27</DATE>

&#5;&#5;&#5;C G
<DATELINE>    LONDON, April 9 - </DATELINE><BODY>Moderately active grain fixing was
reported out of the U.S. But none of the business involved the
significant voyages to the Continent or Japan, ship brokers
    A steady 13.50 dlrs was paid from the U.S. Gulf to Morocco
and 23.25 dlrs was paid for 27,000 long tons from the Gulf to
Taiwan. A vessel carrying 13,500 long tons of bagged wheat
flour from the Gulf to Aqaba received a lump sum of 472,500
    Grain from the Great Lakes to Algeria made 28 dlrs against
27.75 paid for similar fixing towards the end of March.
    Market talk suggested a Federal Commerce vessel had been
booked to move grain from the Great Lakes to Morocco on Comanav
account at about 22 dlrs and 15.50 had been paid for a cargo of
oilseeds from British Columbia to Japan, but no confirmation
was obtainable.
    On the Continent, shippers agreed 19 dlrs for wheat from La
Pallice to Buenaventura and 10.75 dlrs for grain from Ghent to
Naples/Venice range. Elsewhere, maize from East London to Japan
paid 22 dlrs.
    Soviet charterers reappeared in the timecharter sector and
secured a 30,000 tonner from Savona for a trans-Atlantic round
trip at 4,450 dlrs daily and a 31,000 tonner from
Antwerp-Hamburg for a similar voyage at 4,250 dlrs daily.

I've looked for ways to make the code more efficient, but unsuccessfully so. I've run a profiler on my code and it seems like the merge function takes almost all the CPU time during execution. But I don't see any glaring inefficiencies in it.

Is there any way to make my program run faster?

  • \$\begingroup\$ Posted code does not compile. No definition for BigramsMap. \$\endgroup\$ – jwvh Jul 28 at 22:17
  • \$\begingroup\$ Sorry, added it. \$\endgroup\$ – Sandi Jul 29 at 9:07
  • \$\begingroup\$ Can you post the results of your profiler, including the arguments used to run it? \$\endgroup\$ – Mast Jul 29 at 9:09
  • \$\begingroup\$ My guess is that you're paying a toll when you collect the data into SortedMap. As an experiment I replaced every mutable.SortedMap with simply Map. I got what appeared to be the same results (no deep analysis) only the Map elements were in a different order. I would suggest sorting only when and if it's needed. Maybe at the end when the results are organized for presentation. \$\endgroup\$ – jwvh Jul 31 at 7:51
  • \$\begingroup\$ That sounds like a good idea, I will try that. \$\endgroup\$ – Sandi Jul 31 at 14:40

Changing from SortedMap to Map sped up the program massively. One file now only takes a few seconds, instead of over 20 minutes. Sorting can be done at a small cost at bigram querying.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.