Want Help Cracking FAANG?

(Then click this)

×
Back to Question Bank

339. Nested List Weight Sum - Leetcode Solution

Problem Description

The Nested List Weight Sum problem asks you to compute the sum of all integers in a nested list. Each integer is multiplied by its depth in the nested list structure.

You are given a nested list of integers represented by a data structure (often an interface called NestedInteger), where each element is either a single integer or another nested list of integers. The top-level list has depth 1, lists inside it have depth 2, and so on.

Your task is to return the weighted sum of all integers in the list, where each integer is multiplied by the depth at which it is located.

  • Input: A nested list of integers, e.g. [[1,1],2,[1,1]]
  • Output: The sum of all integers in the list, each multiplied by their depth. For the above input, the output is 10.

Key Constraints:

  • Each element is either an integer or a list of elements of the same type.
  • The input list is not empty.
  • There is always one valid solution.

Thought Process

When faced with this problem, the first idea is to simply sum up all the integers. However, the twist is that the sum must be weighted by the depth of each integer. That means, integers deeper inside nested lists contribute more to the final sum.

To solve this, we need to keep track of the current depth as we traverse the nested structure. We can do this recursively: for each element, if it's an integer, multiply it by the current depth; if it's a list, recursively process the list, increasing the depth.

A brute-force approach would be to flatten the list and keep track of the depth for each integer, but recursion or using a stack makes this more natural and efficient.

Optimizing, we realize that recursion neatly handles depth tracking, and there is no need to store all values—just accumulate the weighted sum as we traverse.

Solution Approach

The most effective way to solve this problem is to use a Depth-First Search (DFS) approach, either recursively or with a stack. Here’s how you can approach it:

  1. Define a helper function that takes a nested list and the current depth as arguments.
  2. Iterate through each element in the list:
    • If the element is an integer, multiply it by the current depth and add it to the sum.
    • If the element is a list, recursively call the helper function on that sublist, increasing the depth by 1.
  3. Return the accumulated sum from the helper function.
  4. Start the process by calling the helper function with the original list and a starting depth of 1.

This approach ensures that each integer is counted exactly once, with the correct depth multiplier.

Design Choices:

  • Recursion naturally models the nested structure, making code clear and concise.
  • We avoid extra space by not flattening the list or storing intermediate results.
  • Alternatively, an explicit stack can be used for an iterative solution, but recursion is typically more readable for this problem.

Example Walkthrough

Let's walk through the example input [[1,1],2,[1,1]]:

  • The outermost list is at depth 1. It contains three elements: [1,1], 2, and [1,1].
  • For the first [1,1] (depth 2): each 1 is multiplied by 2, so total is 1*2 + 1*2 = 4.
  • For 2 (depth 1): 2*1 = 2.
  • For the second [1,1] (depth 2): again 1*2 + 1*2 = 4.
  • Add them up: 4 + 2 + 4 = 10.

At each step, the algorithm keeps track of the current depth and multiplies each integer accordingly.

Time and Space Complexity

Brute-force Approach:

  • If you first flatten the list with depth information, then sum up, you traverse all elements twice. Time complexity is O(N), where N is the total number of integers and lists.
  • Space complexity is also O(N) due to storing the flattened structure.
Optimized Recursive Approach:
  • Time complexity is O(N) because each element (integer or list) is visited exactly once.
  • Space complexity is O(D), where D is the maximum depth of nesting, due to the recursion stack.

Both approaches are efficient for reasonable input sizes, but the recursive approach is more space-efficient if the depth is much less than the total number of elements.

Summary

The Nested List Weight Sum problem is elegantly solved by using a recursive approach that tracks the current depth as we traverse the nested structure. By multiplying each integer by its depth, we ensure the correct weighted sum is calculated. Recursion naturally fits the problem, making the solution both simple and efficient. The key insight is to process each element exactly once, leveraging the call stack to manage depth, and thereby achieving optimal time and space complexity.

Code Implementation

# This assumes the NestedInteger interface is provided by Leetcode.
class Solution:
    def depthSum(self, nestedList: List[NestedInteger]) -> int:
        def dfs(nlist, depth):
            total = 0
            for ni in nlist:
                if ni.isInteger():
                    total += ni.getInteger() * depth
                else:
                    total += dfs(ni.getList(), depth + 1)
            return total
        return dfs(nestedList, 1)
      
// This assumes the NestedInteger interface is provided by Leetcode.
class Solution {
public:
    int depthSum(vector<NestedInteger>& nestedList) {
        return dfs(nestedList, 1);
    }
private:
    int dfs(const vector<NestedInteger>& nlist, int depth) {
        int total = 0;
        for (const auto& ni : nlist) {
            if (ni.isInteger()) {
                total += ni.getInteger() * depth;
            } else {
                total += dfs(ni.getList(), depth + 1);
            }
        }
        return total;
    }
};
      
// This assumes the NestedInteger interface is provided by Leetcode.
public class Solution {
    public int depthSum(List<NestedInteger> nestedList) {
        return dfs(nestedList, 1);
    }
    private int dfs(List<NestedInteger> nlist, int depth) {
        int total = 0;
        for (NestedInteger ni : nlist) {
            if (ni.isInteger()) {
                total += ni.getInteger() * depth;
            } else {
                total += dfs(ni.getList(), depth + 1);
            }
        }
        return total;
    }
}
      
// This assumes the NestedInteger interface is provided by Leetcode.
/**
 * @param {NestedInteger[]} nestedList
 * @return {number}
 */
var depthSum = function(nestedList) {
    function dfs(nlist, depth) {
        let total = 0;
        for (let ni of nlist) {
            if (ni.isInteger()) {
                total += ni.getInteger() * depth;
            } else {
                total += dfs(ni.getList(), depth + 1);
            }
        }
        return total;
    }
    return dfs(nestedList, 1);
};