在树中的递归查找

Posted 2023-02-02

我的树形结构如下

  case class Node(val entity : Entity, val children : Seq[Node])
  case class Entity(val key : String)

我在下面编写了根据键过滤节点的函数

  def findPredicate(search: String, entityNode: Node): Boolean =
    search == entityNode.entity.key || entityNode.children.exists(r => findPredicate(search, r))

  def findSubTree(searchKey: String,  entityNode: Node): Option[Node] = 
    val key = entityNode.entity.key
    searchKey match 
      case `key` => Some(entityNode)
      case _ => entityNode.children.find(r => findPredicate(searchKey, r)).
        flatMap(resp => findSubTree(searchKey, resp))
    
  

到目前为止，一切正常，测试用例

  val nodes: Node =
    Node(Entity("o1-01"), Seq(
      Node(Entity("o2-01"), Seq(
        Node(Entity("b-01"), Seq(
          Node(Entity("w-01"), Seq.empty),
          Node(Entity("w-02"), Seq.empty)
        ))
      )),
      Node(Entity("o2-02"), Seq(
        Node(Entity("b-02"), Seq(
          Node(Entity("w-03"), Seq.empty),
          Node(Entity("w-04"), Seq.empty)
        )),
        Node(Entity("b-03"), Seq(
          Node(Entity("w-05"), Seq.empty),
          Node(Entity("w-06"), Seq.empty)
        ))
      ))
    )
    )

  print(findSubTree("o2-01", nodes) + "\n") //Some(Node(Entity(o2-01),List(Node(Entity(b-01),List(Node(Entity(w-01),List()), Node(Entity(w-02),List()))))))
  print(findSubTree("b-01", nodes) + "\n") //Some(Node(Entity(b-01),List(Node(Entity(w-01),List()), Node(Entity(w-02),List()))))
  print(findSubTree("w-01", nodes) + "\n") //Some(Node(Entity(w-01),List()))
  print(findSubTree("w-02", nodes) + "\n") //Some(Node(Entity(w-02),List()))
  print(findSubTree("b-02", nodes) + "\n") //Some(Node(Entity(b-02),List(Node(Entity(w-03),List()), Node(Entity(w-04),List()))))

1. 我想使用多个搜索键和将EntityNode列表作为返回值来扩展findSubTree函数

   def findSubTree（searchKey：Seq [String]，entityNode：Seq [Node]）：选项[Seq [Node]]

2. 我在当前解决方案中已经遍历了2次树。有没有更好的方法基于多个键来查找/过滤节点？

答案

作为深度优先搜索，可以简化当前的findSubTree()。

def findSubTree(searchKey: String,  entityNode: Node): Option[Node] =
  if (searchKey == entityNode.entity.key) Some(entityNode)
  else entityNode.children.view.flatMap(findSubTree(searchKey, _)).headOption

可以利用此设计以适合增强的findSubTrees()（注意，我更改了名称）。

def findSubTrees(searchKeys: Seq[String], entityNodes: Seq[Node]): Option[Seq[Node]] = 
  def findNode(node :Node) :Option[Node] =
    if (searchKeys contains node.entity.key) Some(node)
    else node.children.view.flatMap(findNode).headOption

  val res = entityNodes.flatMap(findNode)
  if (res.isEmpty) None
  else Some(res)

另一答案

这里是使用Stream的方法。递归直到找到匹配项。

Stream

样本Scastie：def findSub(n: Node, k: String): Stream[Node] = if (n.entity.key == k) Stream(n) else n.children.toStream.flatMap( findSub(_, k) )