[ネタ] Trie を実装してみる

Trie.scala

// mutable な trie と immutable な trie
// 久しぶりに scala 書いてみたらダーティ過ぎてワロタ。。。
// 仕様は http://d.hatena.ne.jp/route150/20120624/1340542890
// 破壊的操作はいいぞ（以下略

package trie {
  // trie base 
  abstract class Trie[A] {
    type ChildrenT <: collection.Map[A, Trie[A]]
    val children: ChildrenT
    
    def add(a: Seq[A]): Trie[A]
    
    def find(a: Seq[A]): Option[Trie[A]] = a match {
      case Seq() =>
        Some(this)
      case Seq(x, xs@_*) =>
        children.get(x).flatMap(_.find(xs))
    }
    
    def allPostfixes: Stream[Stream[A]]
    
    protected final def dfsPostfixes: Stream[Stream[A]] = for {
      (k, v) <- children.iterator.toStream
      pf <- v.allPostfixes
    } yield k #:: pf
    
    def autoComplete(xs: Seq[A]): Stream[Stream[A]] =
      find(xs).map{_.allPostfixes}.getOrElse(Stream.empty)   
  }
  
  // immutable trie
  package immutable {
    abstract class Trie[A] extends trie.Trie[A] {
      type ChildrenT = Map[A, Trie[A]]
      
      def copyWithChildren(children: Map[A, Trie[A]]): Trie[A]
      
      def add(a: Seq[A]): Trie[A] = a match {
        case Seq() => Trie.Leaf(children)
        case Seq(x, xs@_*) =>
          copyWithChildren(children + (x -> children.getOrElse(x, Trie()).add(xs)))
      }
    }
    
    object Trie {
      def apply[A](): Trie[A] = Node[A](Map.empty)
      
      case class Node[A](val children: Map[A, Trie[A]]) extends Trie[A] {
        def allPostfixes = dfsPostfixes
        def copyWithChildren(c: Map[A, Trie[A]]) = copy(c)
      }
      
      case class Leaf[A](val children: Map[A, Trie[A]]) extends Trie[A] {
        def allPostfixes = Stream.empty #:: dfsPostfixes
        def copyWithChildren(c: Map[A, Trie[A]]) = copy(c)
      }
    }  
  }
  
  // mutable trie
  package mutable {
    import scala.collection.mutable.Map
    
    case class Trie[A](
        children: Map[A, Trie[A]] = Map.empty[A, Trie[A]]
        ) extends trie.Trie[A] {
      
      type ChildrenT = Map[A, Trie[A]]
      var flag = false
      
      def add(a: Seq[A]): Trie[A] = {
        a match {
          case Seq() =>
            flag = true
          case Seq(x, xs@_*) =>
            children += (x -> children.getOrElse(x, Trie()).add(xs))
        }
        this
      }
      
      def allPostfixes =
        if (flag)
          Stream.Empty #:: dfsPostfixes
        else
          dfsPostfixes
    }
  }
}

object Main extends App {
  val x = trie.mutable.Trie[Char]()
  x.add("hoge")
  x.add("hage")
  println(x.allPostfixes.map{_.mkString}.toList)
  println(x.autoComplete("hog").map{_.mkString}.toList)
  println(x.toString)
  
  val a = trie.immutable.Trie[Int]()
  val b = a.add(Vector(1, 2, 3, 4))
  val c = b.add(List(1, 2, 5, 6))
  println(c.allPostfixes.map(_.force).force)
  println(b.allPostfixes.map(_.force).force)
  println(c.toString)
}