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]], return10. (four 1's at depth 2, one 2 at depth 1)

Example 2: Given the list[1,[4,[6]]], return27. (one 1 at depth 1, one 4 at depth 2, and one 6 at depth 3; 1 + 4*2 + 6*3 = 27)

# """
# This is the interface that allows for creating nested lists.
# You should not implement it, or speculate about its implementation
# """
#class NestedInteger(object):
#    def isInteger(self):
#        """
#        @return True if this NestedInteger holds a single integer, rather than a nested list.
#        :rtype bool
#        """
#
#    def getInteger(self):
#        """
#        @return the single integer that this NestedInteger holds, if it holds a single integer
#        Return None if this NestedInteger holds a nested list
#        :rtype int
#        """
#
#    def getList(self):
#        """
#        @return the nested list that this NestedInteger holds, if it holds a nested list
#        Return None if this NestedInteger holds a single integer
#        :rtype List[NestedInteger]
#        """

class Solution(object):
    def depthSum(self, nestedList):
        """
        :type nestedList: List[NestedInteger]
        :rtype: int
        """
        return self.depthSumHelper(nestedList, 1)
        #sum = self.depthSumHelper(nestedList, 1)
        #print 'final sum: ', sum
        #return sum

    def depthSumHelper(self, nestedList, rep, sum = 0): #rep stands for repetition/or depth of the DFS
        for item in nestedList:
            #print 'item'
            if item.getList():
                sum = self.depthSumHelper(item.getList(), rep + 1, sum)
                #print 'retSum: ', sum
            elif item.isInteger():
                i = item.getInteger() * rep
                sum += i
                #print 'integer: ', i, 'rep: ', rep, 'sum: ', sum
        #print 'returning: ', sum
        return sum

There is a nice editorial discussion on time/space complexity:

Ref: https://leetcode.com/articles/nested-list-weight-sum/

Complexity Analysis

The algorithm takes O(N) time, where N is the total number of nested elements in the input list. For example, the list[ [[[[1]]]], 2 ]contains 4 nested lists and 2 nested integers (1 and 2), so N=6

In terms of space, at most O(D) recursive calls are placed on the stack, where D is the maximum level of nesting in the input. For example, D=2 for the input[[1,1],2,[1,1]], and D=3 for the input[1,[4,[6]]].

Last updated

Was this helpful?