Nested List Weight Sum

Given a nested list of integers, return the sum of all integers in the list weighted by their depth.

Each element is either an integer, or a list -- whose elements may also be integers or other lists.

Example 1: Given the list

[[1,1],2,[1,1]]

return

10 (four 1's at depth 2, one 2 at depth 1)

Example 2: Given the list

[1,[4,[6]]]

return

27 (one 1 at depth 1, one 4 at depth 2, and one 6 at depth 3; 1 + 42 + 63 = 27)

trait NestedItem {
  def getInt: Int

  def getList: List[NestedItem]
}

class NestedInt(value: Int) extends NestedItem {
  def getInt: Int = value

  def getList: List[NestedItem] = throw new Error("get list from int")
}

class NestedList(value: List[NestedItem]) extends NestedItem {
  def getInt: Int = throw new Error("get int from list")

  def getList = value
}

def depthSum(nestedList: List[NestedItem], depth: Int = 1): Int = {
  var sum = 0
  for (item <- nestedList) {
    item match {
      case i: NestedInt => sum += i.getInt * depth
      case lst: NestedList => sum += depthSum(lst.getList, depth + 1)
    }
  }
  sum
}

def wrap(lst: List[Any]): NestedList = {
  new NestedList(lst.map((item) => item match {
    case i: Int => new NestedInt(i)
    case l: List[Any] => wrap(l)
  }))
}

val test1 = wrap(List(List(1, 1), 2, List(1, 1)))
depthSum(test1.getList)

val test2 = wrap(List(1, List(4, List(6))))
depthSum(test2.getList)

Last updated