LeetCode 364: Nested List Weight Sum II

link

Two pass

Forest

Recursive

Each NestedInteger can be thought of as a tree. So, the problem involves a forest of trees. We first find the maximum height of these trees which is the max depth for the forest. Then, we sum up the weighted sum of all trees in the forest.

Time: \mathcal{O}(n), space: \mathcal{O}(\texttt{forest-height}).

# """
# This is the interface that allows for creating nested lists.
# You should not implement it, or speculate about its implementation
# """
# class NestedInteger:
#    def __init__(self, value=None):
#        """
#        If value is not specified, initializes an empty list.
#        Otherwise initializes a single integer equal to value.
#        """
#
#    def isInteger(self):
#        """
#        @return True if this NestedInteger holds a single integer, rather than a nested list.
#        :rtype bool
#        """
#
#    def add(self, elem):
#        """
#        Set this NestedInteger to hold a nested list and adds a nested integer elem to it.
#        :rtype void
#        """
#
#    def setInteger(self, value):
#        """
#        Set this NestedInteger to hold a single integer equal to value.
#        :rtype void
#        """
#
#    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:
    def height(self, root) -> int:
        if root.isInteger():
            return 0
        return max(self.height(c) for c in root.getList()) + 1

    def depthSumInverse(self, nestedList: List[NestedInteger]) -> int:
        max_depth = max(self.height(root) for root in nestedList)

        def weighted_sum(root, depth) -> int:
            nonlocal max_depth

            if root.isInteger():
                return (max_depth - depth + 1) * root.getInteger()

            return sum(weighted_sum(c, depth + 1) for c in root.getList())

        return sum(weighted_sum(root, 0) for root in nestedList)

Tree

Recursive

We can create a top-level NestedInteger that ties the trees of the forest. In that way we have a single tree to work with. Note, depth of each leaf would be one more than the forest approach. However, the weight, (max_depth - depth + 1) will remain the same.

# """
# This is the interface that allows for creating nested lists.
# You should not implement it, or speculate about its implementation
# """
# class NestedInteger:
#    def __init__(self, value=None):
#        """
#        If value is not specified, initializes an empty list.
#        Otherwise initializes a single integer equal to value.
#        """
#
#    def isInteger(self):
#        """
#        @return True if this NestedInteger holds a single integer, rather than a nested list.
#        :rtype bool
#        """
#
#    def add(self, elem):
#        """
#        Set this NestedInteger to hold a nested list and adds a nested integer elem to it.
#        :rtype void
#        """
#
#    def setInteger(self, value):
#        """
#        Set this NestedInteger to hold a single integer equal to value.
#        :rtype void
#        """
#
#    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:
    def height(self, root: NestedInteger) -> int:
        if root.isInteger():
            return 0
        
        return max( self.height(c) for c in root.getList() ) + 1

    def depthSumInverse(self, nestedList: List[NestedInteger]) -> int:
        root = NestedInteger()
        for c in nestedList:
            root.add(c)

        max_depth = self.height(root)

        def weighted_sum(root: NestedInteger, depth: int) -> int:
            nonlocal max_depth

            if root.isInteger():
                return (max_depth-depth+1) * root.getInteger()
            
            return sum( weighted_sum(c, depth+1) for c in root.getList() )

        return weighted_sum(root, 0)

Iterative

We simulate the tree approach. There are two bottom-up recursions: (1) height (2) weighted_sum. To simulate each recursion, we use a stack, a visited set, and a node_val dict.

# """
# This is the interface that allows for creating nested lists.
# You should not implement it, or speculate about its implementation
# """
# class NestedInteger:
#    def __init__(self, value=None):
#        """
#        If value is not specified, initializes an empty list.
#        Otherwise initializes a single integer equal to value.
#        """
#
#    def isInteger(self):
#        """
#        @return True if this NestedInteger holds a single integer, rather than a nested list.
#        :rtype bool
#        """
#
#    def add(self, elem):
#        """
#        Set this NestedInteger to hold a nested list and adds a nested integer elem to it.
#        :rtype void
#        """
#
#    def setInteger(self, value):
#        """
#        Set this NestedInteger to hold a single integer equal to value.
#        :rtype void
#        """
#
#    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:
    def height(self, root: NestedInteger) -> int:
        visited, node_height = set(), {}
        stack = [root]
        while stack:
            u = stack.pop()
            if u.isInteger():
                node_height[u] = 0
                continue
            if u not in visited:
                visited.add(u)
                stack.append(u)
                stack.extend(c for c in u.getList())
                continue
            # Return path: heights of u's children are available
            node_height[u] = max( node_height[c] for c in u.getList() ) + 1

        return node_height[root]


    def depthSumInverse(self, nestedList: List[NestedInteger]) -> int:
        root = NestedInteger()
        for c in nestedList:
            root.add(c)

        max_depth = self.height(root)

        visited, node_weighted_sum = set(), {}
        stack = [(root, 0)]
        while stack:
            u, depth = stack.pop()
            if u.isInteger():
                node_weighted_sum[u] = (max_depth-depth+1) * u.getInteger()
                continue
            if u not in visited:
                visited.add(u)
                stack.append((u, depth))
                stack.extend( (c, depth+1) for c in u.getList() )
                continue
            # Return path: weighted sums of u's children are available
            node_weighted_sum[u] = sum( node_weighted_sum[c] for c in u.getList() )

        return node_weighted_sum[root]

One pass

Tree

Recursive

\texttt{Sum}_\texttt{tree} = \sum_{u \in \texttt{tree}} \left( \texttt{depth}_{max} - \texttt{depth}_{u} + 1 \right) \cdot \texttt{val}_{u} = \left( \texttt{depth}_{max} + 1 \right) \ \sum_{u} \texttt{val}_{u} - \sum_{u} \texttt{depth}_{u} \cdot \texttt{val}_{u}.

So, with a single DFS we can find \texttt{depth}_{max}, \sum_{u} \texttt{val}_{u}, and \sum_{u} \texttt{depth}_{u} \cdot \texttt{val}_{u}.

# """
# This is the interface that allows for creating nested lists.
# You should not implement it, or speculate about its implementation
# """
# class NestedInteger:
#    def __init__(self, value=None):
#        """
#        If value is not specified, initializes an empty list.
#        Otherwise initializes a single integer equal to value.
#        """
#
#    def isInteger(self):
#        """
#        @return True if this NestedInteger holds a single integer, rather than a nested list.
#        :rtype bool
#        """
#
#    def add(self, elem):
#        """
#        Set this NestedInteger to hold a nested list and adds a nested integer elem to it.
#        :rtype void
#        """
#
#    def setInteger(self, value):
#        """
#        Set this NestedInteger to hold a single integer equal to value.
#        :rtype void
#        """
#
#    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:

    def depthSumInverse(self, nestedList: List[NestedInteger]) -> int:
        root = NestedInteger()
        for c in nestedList:
            root.add(c)

        max_depth, plain_sum, weighted_sum = 0, 0, 0
        def dfs(root, depth) -> None:
            nonlocal max_depth, plain_sum, weighted_sum

            if root.isInteger():
                max_depth = max(max_depth, depth)
                val = root.getInteger()
                plain_sum += val
                weighted_sum += (depth * val)
                return
            
            for c in root.getList():
                dfs(c, depth+1)
        
        dfs(root, 0)
        return (max_depth+1)*plain_sum - weighted_sum

Iterative

We use an explicit stack to simulate the bottom-up recursion.

# """
# This is the interface that allows for creating nested lists.
# You should not implement it, or speculate about its implementation
# """
# class NestedInteger:
#    def __init__(self, value=None):
#        """
#        If value is not specified, initializes an empty list.
#        Otherwise initializes a single integer equal to value.
#        """
#
#    def isInteger(self):
#        """
#        @return True if this NestedInteger holds a single integer, rather than a nested list.
#        :rtype bool
#        """
#
#    def add(self, elem):
#        """
#        Set this NestedInteger to hold a nested list and adds a nested integer elem to it.
#        :rtype void
#        """
#
#    def setInteger(self, value):
#        """
#        Set this NestedInteger to hold a single integer equal to value.
#        :rtype void
#        """
#
#    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:

    def depthSumInverse(self, nestedList: List[NestedInteger]) -> int:
        root = NestedInteger()
        for c in nestedList:
            root.add(c)

        max_depth, plain_sum, weighted_sum = 0, 0, 0
        stack = [(root, 0)]
        while stack:
            u, depth = stack.pop()
            if u.isInteger():
                max_depth = max( max_depth, depth )
                val = u.getInteger()
                plain_sum += val
                weighted_sum += (depth * val)
                continue
            
            stack.extend((c, depth+1) for c in u.getList())

        return (max_depth+1)*plain_sum - weighted_sum

Leave a comment