Playing with trees and distances between nodes / published by https://github.com/dacr/code-examples-manager #8b946284-42d7-4aa7-bf1e-5758be2783b4/202044a3e1e89f2cecd1cf683832930e5bcfb564

algo-trees.sc

// summary : Playing with trees and distances between nodes
// keywords : scala, algorithm, scalatest, tree, @testable
// publish : gist
// authors : David Crosson
// license : Apache NON-AI License Version 2.0 (https://raw.githubusercontent.com/non-ai-licenses/non-ai-licenses/main/NON-AI-APACHE2)
// id : 8b946284-42d7-4aa7-bf1e-5758be2783b4
// created-on : 2019-12-06T22:06:08Z
// managed-by : https://github.com/dacr/code-examples-manager
// run-with : scala-cli $file

// ---------------------
//> using scala "3.4.2"
//> using dep "org.scalatest::scalatest:3.2.16"
//> using objectWrapper
// ---------------------

import org.scalatest._, flatspec.AnyFlatSpec, matchers.should.Matchers
import org.scalatest.OptionValues._

case class Node(name: String)
case class Tree(node: Node, children: List[Tree]=Nil)

def buildTree(treeDesc: String, connectionMarkerRegex:String="[)]"): List[Tree] = {
  val connections =
    treeDesc
      .split("\n")
      .map(_.split(connectionMarkerRegex, 2))
      .map { case Array(aName, bName) => Node(aName) -> Node(bName) }
      .toList
      .groupMap { case (a, _) => a } { case (_, b) => b }
  val children = connections.values.flatten.toList
  val nodes = connections.keys.toSet ++ children
  val rootNodes = nodes -- children
  def build(fromNode: Node, remainingConnections: Map[Node, List[Node]]): Tree = {
    val newRemainingConnections = remainingConnections - fromNode
    val subtrees =
      remainingConnections
        .get(fromNode)
        .getOrElse(Nil)
        .map(node => build(node,newRemainingConnections))
    Tree(fromNode, subtrees)
  }
  rootNodes.toList.map(rootNode => build(rootNode, connections))
}


// distance from root to the given node
def distanceTo(tree:Tree, node:Node):Option[Int] = {
  def searchDistance(current:Tree, node:Node, depth:Int):Option[Int] = {
    if (current.node == node) Some(depth)
    else if (current.children.isEmpty) None
    else {
      current
        .children
        .to(LazyList)
        .map(subTree => searchDistance(subTree, node, depth+1))
        .find(_.isDefined)
        .flatten
    }
  }
  searchDistance(tree, node, 0)
}

// minimal distance between two nodes
def distanceBetween(tree: Tree, nodeA: Node, nodeB: Node): Option[Int] = {
  def explore(currentTree:Tree, bestDistance:Option[Int]):Option[Int] = {
    (distanceTo(currentTree, nodeA), distanceTo(currentTree, nodeB)) match {
      case (None, _) => bestDistance
      case (_, None) => bestDistance
      case (Some(da), Some(db)) =>
        val newBestDistance = bestDistance.filter(_ < (da+db)).orElse(Some(da+db))
        val results = currentTree.children.map(child => explore(child, newBestDistance)).flatten
        results.minByOption(x => x).orElse(newBestDistance)
    }
  }
  explore(tree, None)
}



class TreeTest extends AnyFlatSpec with Matchers {
  override def suiteName: String = "TreeTest"
  "buildTree" should "be able to build a one root/leaf tree" in {
    val treeDesc = "a)b"
    buildTree(treeDesc).head shouldBe Tree(Node("a"),List(Tree(Node("b"),Nil)))
  }
  it should "be able to build a one root, two leaves tree" in {
    val treeDesc = "a)b\na)c"
    buildTree(treeDesc).head shouldBe Tree(Node("a"),List(Tree(Node("b")), Tree(Node("c"))))
  }
  it should "be able several trees" in {
    val treeDesc = "a)b\nc)d"
    buildTree(treeDesc) should contain allOf (
      Tree(Node("a"),List(Tree(Node("b")))),
      Tree(Node("c"),List(Tree(Node("d")))),
    )
  }
  "distanceTo" should "return the distance from root to a give node" in {
    val treeDesc = "a)b\nb)c"
    val tree = buildTree(treeDesc).headOption.value
    distanceTo(tree, Node("a")).value shouldBe 0
    distanceTo(tree, Node("b")).value shouldBe 1
    distanceTo(tree, Node("c")).value shouldBe 2
  }
  "distanceBetween" should "return the shortest distance between two nodes" in {
    val treeDesc = "root)a\na)b\nb)c\nb)d\nd)e\ne)f"
    val tree = buildTree(treeDesc).headOption.value
    distanceBetween(tree, Node("c"), Node("f")).value shouldBe 4
  }
}

org.scalatest.tools.Runner.main(Array("-oDF", "-s", classOf[TreeTest].getName))