Search code examples
scalasizefoldleft

How to get a size of a tree-like custom object


I am trying to figure out how to get the size (num of "levels") of a custom tree-like data structure defined this way:

case class PrefixTreeImpl[K, +V](value: Option[V], prefixTree: Map[K, PrefixTree[K, V]]) extends PrefixTree[K, V]

As you can see it implements the PrefixTree interface (trait actually) which has a size() method amongst others:

def size: Int

I am trying to implement it using foldLeft() Scala method, but I don't really understand how it works with this kind of sophisticated data structures So far I came up with this and got stuck:

override def size: Int = {
  if (prefixTree.isEmpty) 0
  else
    prefixTree.foldLeft(0) {(z, Map[K, PrefixTree[K, V]]()) => z + 1}

}

Obviously it doesn't compile, but I couldn't come up with something else yet.

The way the DS works is:

in: scala>val tree2 = tree1.put(List("one", "two"), 12)

out: PrefixTreeImpl(None,Map(one -> PrefixTreeImpl(None,Map(two -> PrefixTreeImpl(Some(12),Map(tree -> PrefixTreeImpl(Some(123),Map())))))))


Solution

  • It can be done recursively in a next way:

    override def size: Int = {
      if(prefixTree.nonEmpty) prefixTree.values.map(_.size + 1).max else 0
    }
    

    So next test:

    val tree = PrefixTreeImpl(None ,Map("one" -> PrefixTreeImpl(None, Map("two" -> PrefixTreeImpl(Some(12),Map("tree" -> PrefixTreeImpl(Some(123), Map.empty[String, PrefixTree[String, Int]])))))))
        println(tree.size)
    

    produces result: 3

    Hope this helps!