Scala Trie

Trie.scala

case class Trie(char: Char,
                children: Seq[Trie] = Nil,
                indices: Seq[Int] = Nil) {

  def add(word: String,
          index: Int): Trie = {
    val updatedChildren = word.toSeq match {
      case Seq(head) =>
        update(head) {
          Trie(head, indices = Seq(index))
        } { trie =>
          trie.copy(indices = trie.indices :+ index)
        }
      case head +: tail =>
        update(head) {
          Trie(head).add(tail.mkString, index)
        } { trie =>
          trie.add(tail.mkString, index)
        }
    }

    copy(children = updatedChildren)
  }

  private def update(head: Char)
                    (notFound: => Trie)
                    (found: Trie => Trie) =
    children.partition(_.char == head) match {
      case (Nil, tries) => notFound +: tries
      case (trie +: _, tries) => found(trie) +: tries
    }

  def get(word: String): Seq[Int] =
    word.toSeq match {
      case Seq(head) =>
        children.find(_.char == head).map(_.indices).getOrElse(Nil)
      case head +: tail =>
        children.find(_.char == head).map(_.get(tail.mkString)).getOrElse(Nil)
    }
}

assert(
  Trie(' ')
    .add("hello", 0)
    .add("world", 1)
    .add("hello", 2)
    .get("hello") == Seq(0, 2)
)