Problem

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.

Examples

Example 1:

Input: [ [1,1],2,[1,1] ]
Output: 10 
Explanation: Four 1's at depth 2, one 2 at depth 1.
elements11211Sum of elements * depth
depths221221x2+1x2+2x1+1x2+1x2 = 10

Example 2:

Input: [1,[ 4,[ 6 ] ] ]
Output: 27 
Explanation: One 1 at depth 1, one 4 at depth 2, and one 6 at depth 3; 1 + 4*2 + 6*3 = 27.
elements146Sum of elements * depth
depths1231x1 + 4x2 + 6x3 = 27

Solution

Method 1 - Recursive

Here is the approach:

  1. Depth-First Search (DFS): Use a recursive approach to traverse the nested list.
  2. Calculate Weighted Sum: For each integer, compute the weighted sum by multiplying the integer with its depth.
  3. Recurse into Sub-Lists: If an element is a list, recursively compute its weighted sum by increasing the depth.

Code

Java

Here is NestedInteger interface:

// This is the interface that allows for creating nested lists.
// You should not implement it, or speculate about its implementation
public interface NestedInteger {
    // Constructor initializes an empty nested list.
    public NestedInteger();

    // Constructor initializes a single integer.
    public NestedInteger(int value);

    // @return true if this NestedInteger holds a single integer, rather than a
    // nested list.
    public boolean isInteger();

    // @return the single integer that this NestedInteger holds, if it holds a
    // single integer Return null if this NestedInteger holds a nested list
    public Integer getInteger();

    // Set this NestedInteger to hold a single integer.
    public void setInteger(int value);

    // Set this NestedInteger to hold a nested list and adds a nested integer to
    // it.
    public void add(NestedInteger ni);

    // @return the nested list that this NestedInteger holds, if it holds a
    // nested list
    // Return empty list if this NestedInteger holds a single integer
    public List<NestedInteger> getList();
}

class Solution {
    public int depthSum(List<NestedInteger> nestedList) {
        return dfs(nestedList, 1);
    }

    private int dfs(List<NestedInteger> nestedList, int depth) {
		int total = 0;
		for (NestedInteger ni : nestedList) {
			if (ni.isInteger()) {
				total += ni.getInteger() * depth;
			} else {
				total += depthSumHelper(ni.getList(), depth + 1);
			}
		}
		return total;
    }
}
Python

Here is NestedInteger class:

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 depthSum(self, nestedList: List[NestedInteger]) -> int:
        def dfs(nestedList, depth):
            depth_sum = 0
            for item in nestedList:
                if item.isInteger():
                    depth_sum += item.getInteger() * depth
                else:
                    depth_sum += dfs(item.getList(), depth + 1)
            return depth_sum

        return dfs(nestedList, 1)

Complexity

  • ⏰ Time complexity: O(n), where n is the total number of elements (both integers and lists) in the nested list. Each element is visited once.
  • 🧺 Space complexity: O(d), where d is the maximum depth of the nested list. This accounts for the space used by the recursion stack.

Method 2 - Iterative

Here are the steps:

  1. Initialize Stack:
    • Use a stack to keep track of the current list and its depth. The stack is initialized with the input list and depth 1.
  2. Iterate Through Stack:
    • While the stack is not empty, pop elements and process them.
    • For each element in the current list:
      • If it is an integer, add its weighted value (element * depth) to the total sum.
      • If it is a sublist, push the sublist onto the stack with an increased depth (depth + 1).
  3. Return Result:
    • After processing all elements, return the total weighted sum.

Code

Java
public class Solution {
    public int depthSum(List<NestedInteger> nestedList) {
        int total = 0;
        Stack<Pair<List<NestedInteger>, Integer>> stack = new Stack<>();
        stack.push(new Pair<>(nestedList, 1));

        while (!stack.isEmpty()) {
            Pair<List<NestedInteger>, Integer> current = stack.pop();
            List<NestedInteger> currentList = current.getKey();
            int depth = current.getValue();

            for (NestedInteger ni : currentList) {
                if (ni.isInteger()) {
                    total += ni.getInteger() * depth;
                } else {
                    stack.push(new Pair<>(ni.getList(), depth + 1));
                }
            }
        }

        return total;
    }
}
Python
class Solution:
    def depthSum(self, nestedList: List[NestedInteger]) -> int:
        total = 0
        stack = [
            (nestedList, 1)
        ]  # Each element is a tuple (current_list, current_depth)

        while stack:
            current_list, depth = stack.pop()
            for item in current_list:
                if item.isInteger():
                    total += item * depth  # Add weighted value
                else:
                    stack.append(
                        (item, depth + 1)
                    )  # Add nested list with increased depth

        return total

Complexity

  • ⏰ Time complexity: O(n), where n is the total number of elements (both integers and lists) in the nested list. Each item is visited once.
  • 🧺 Space complexity: O(n), where n is the total number of items in the nested list. This accounts for the space used by the stack.