木構造の再帰による深さ優先検索

木構造の再帰を使った巡回 という エントリーのコードの改良版です。

なぜか最近たびたび木構造を扱うことがあったので、その辺を整理を含めた覚え書きです。

環境

$ kotlin -version
Kotlin version 1.9.22-release-704 (JRE 17.0.10+7-Ubuntu-122.04.1)

木構造のノードを表現する Tree クラス

以前のエントリーでは、木を次のように表現していました。

data class Node(val id: String, val children: List<Node> = listOf())

別にこれでも機能するのですが、今回はこのように表現することにします。

typealias ID = String

sealed class Tree(open val id: ID) {
    data class Leaf(override val id: ID): Tree(id)
    data class Node(override val id: ID, val children: List<Tree>): Tree(id)
}

このように 木を sealed クラスで表現した方がこのクラスを使う側のコードが書きやすい気がする。

木構造の定義

テストに使う木構造を定義。

val rootTree =
  Tree.Node("0", listOf(
    Tree.Node("01", listOf(
      Tree.Leaf("01-01"),
      Tree.Leaf("01-02"),
      Tree.Leaf("01-03"))),
    Tree.Node("02", listOf(
      Tree.Leaf("02-01"),
      Tree.Leaf("02-02"),
      Tree.Leaf("02-03"))),
    Tree.Node("03", listOf(
      Tree.Leaf("03-01"),
      Tree.Leaf("03-02"),
      Tree.Leaf("03-03")))))

木構造を深さ優先で探索する関数 depthFirstSearch を使ってフラットにする

Tree を List<ID> にする。（つまりフラットにする）

typealias Converter: (Tree)->List<ID>

肝心の実装:

val depthFirstSearch: Converter = { tree->

    fun f(t: Tree): List<String> {
        return when( t ) {
            is Tree.Leaf -> { listOf(t.id) }
            is Tree.Node -> {
                t.children.fold(listOf(t.id), { acc, childT-> acc + f(childT) })
                //listOf(t.id) + t.children.map { childT-> f(childT) }.flatten()
            }
        }
    }

    f(tree)
}

実行して確かめる:

val idList: List<ID> = depthFirstSearch(rootTree)

idList.forEach {
    val indent = 0.until(it.id.length).map { " " }.joinToString("")
    println("${indent}- ${it.id}")
}

結果が分かりやすいようにインデントを計算しています。 id の文字列の長さから深さをきめるという適当なコード。

$ kotlin main.kts
 - 0
  - 01
     - 01-01
     - 01-02
     - 01-03
  - 02
     - 02-01
     - 02-02
     - 02-03
  - 03
     - 03-01
     - 03-02
     - 03-03

できた。

まとめ

コード全体を掲載。

main.kts

typealias ID = String
typealias Converter = (Tree)->List<ID>

sealed class Tree(open val id: ID) {
    data class Leaf(override val id: ID): Tree(id)
    data class Node(override val id: ID, val children: List<Tree>): Tree(id)
}

val depthFirstSearch: Converter = { tree->

    fun f(t: Tree): List<String> {
        return when( t ) {
            is Tree.Leaf -> { listOf(t.id) }
            is Tree.Node -> {
                t.children.fold(listOf(t.id), { acc, childT-> acc + f(childT) })
                //listOf(t.id) + t.children.map { childT-> f(childT) }.flatten()
            }
        }
    }

    f(tree)
}

val rootTree =
  Tree.Node("0", listOf(
    Tree.Node("01", listOf(
      Tree.Leaf("01-01"),
      Tree.Leaf("01-02"),
      Tree.Leaf("01-03"))),
    Tree.Node("02", listOf(
      Tree.Leaf("02-01"),
      Tree.Leaf("02-02"),
      Tree.Leaf("02-03"))),
    Tree.Node("03", listOf(
      Tree.Leaf("03-01"),
      Tree.Leaf("03-02"),
      Tree.Leaf("03-03")))))

val idList: List<ID> = depthFirstSearch(rootTree)

idList.forEach { id->
    val indent = 0.until(id.length).map { " " }.joinToString("")
    println("${indent}- ${id}")
}

次のようにして実行。

$ kotlin main.kts

以上です。