多分木をフラットにする, Scala

以前メモ的に書いたこれ、少し使い勝手が悪いのと、この手の処理はぱっと書けないのでメモする。

tchiba.hatenablog.jp

Node trait

上の記事で作った、単純な多分木をフラットにする処理を拡張してみた。

import scala.annotation.tailrec

trait Node[A <: Node[_]]{
  val children: List[A]

  protected def flatChildren[B](z: B)(f: A => B): List[B] = { // 継承先の具体的な実装での利用を想定しているので、protected にした。
    @tailrec
    def loop(children: List[A], vs: List[B]): List[B] = {
      children match {
        case ::(head, next) => loop(next ++ head.children.asInstanceOf[List[A]], vs.appended(f(head)))
        case Nil => vs
      }
    }

    loop(children, List(z))
  }
}

使い方

Node trait を継承して、利用先の具象クラスの振る舞いにふさわしい表現のメソッド名をつけてあげてください。

case class SampleTree(
    value: Int,
    children: List[SampleTree]
) extends Node[SampleTree] {

  // 第一引数には初期値を、第二引数には元となる木をベースにした変換処理を書きます。
  def allValues: List[Int] = flatChildren(value)(t => t.value)
}

実行

// サンプルデータ
val tree = {
  SampleTree(
    1,
    List(
      SampleTree(2, List.empty),
      SampleTree(3, List.empty),
      SampleTree(4, List(
        SampleTree(5, List.empty),
        SampleTree(6, List.empty),
        SampleTree(7, List(
          SampleTree(8, List(
            SampleTree(9, List(
              SampleTree(12, List(
                SampleTree(13, List.empty)
              ))
            ))
          )),
          SampleTree(10, List(
            SampleTree(11, List.empty)
          ))
        ))
      ))
    )
  )
}

// 確認
println(tree.allValues)
// List(1, 2, 3, 4, 5, 6, 7, 8, 10, 9, 11, 12, 13)

木がおきくなると計算量も多くなるので、lazy val とかにしておいた方がいいもしれません。

case class SampleTree(
    value: Int,
    children: List[SampleTree]
) extends Node[SampleTree] {

  lazy val allValues: List[Int] = flatChildren(value)(t => t.value)
}