summaryrefslogtreecommitdiff
path: root/core/source/core/scala/com/rockymadden/stringmetric/tokenization/NGramTokenizer.scala
blob: d66fd62fbf8627e38b7fbbc5d4b012efc108fa82 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
package com.rockymadden.stringmetric.tokenization

import com.rockymadden.stringmetric.{StringFilter, StringTokenizer}
import scala.annotation.tailrec

/** An implementation of the N-Gram tokenizer. */
class NGramTokenizer extends StringTokenizer[Int, Array[String]] { this: StringFilter =>
	final override def tokenize(charArray: Array[Char])(implicit n: Int): Option[Array[Array[Char]]] = {
		if (n <= 0) throw new IllegalArgumentException("Expected valid n.")

		val fca = filter(charArray)

		if (fca.length < n) None
		else Some(sequence(fca, Array.empty[Array[Char]], n))
	}

	final override def tokenize(string: String)(implicit n: Int): Option[Array[String]] =
		tokenize(string.toCharArray)(n).map(_.map(_.mkString))

	@tailrec
	private[this] def sequence(i: Array[Char], o: Array[Array[Char]], n: Int): Array[Array[Char]] = {
		require(n > 0)

		if (i.length <= n) o :+ i
		else sequence(i.tail, o :+ i.take(n), n)
	}
}

object NGramTokenizer {
	private lazy val self = apply()

	def apply(): NGramTokenizer = new NGramTokenizer with StringFilter

	def tokenize(charArray: Array[Char])(n: Int) = self.tokenize(charArray)(n)

	def tokenize(string: String)(n: Int) = self.tokenize(string)(n)
}