LeetCode 339: Nested List Weight Sum

Problem Restatement

We are given a nested list of integers called nestedList.

Each element is either:

The depth of an integer is the number of lists it is inside.

Return the sum of every integer multiplied by its depth. The problem examples include [[1,1],2,[1,1]] -> 10, [1,[4,[6]]] -> 27, and [0] -> 0.

Input and Output

Function shape:

def depthSum(nestedList: list[NestedInteger]) -> int:
    ...

Examples

Example 1:

Input: nestedList = [[1,1],2,[1,1]]
Output: 10

The four 1 values are at depth 2.

The value 2 is at depth 1.

So:

1 * 2 + 1 * 2 + 2 * 1 + 1 * 2 + 1 * 2 = 10

Example 2:

Input: nestedList = [1,[4,[6]]]
Output: 27

The values are:

Total:

1 + 8 + 18 = 27

Example 3:

Input: nestedList = [0]
Output: 0

The only integer is 0, so the weighted sum is 0.

First Thought: Traverse Everything

We need to visit every integer exactly once.

The only extra information needed during traversal is the current depth.

This naturally suggests recursion:

1. Start at depth 1.
2. For each element:
   1. If it is an integer, add value * depth.
   2. If it is a list, recursively process it at depth + 1.

This is depth-first search over a tree-like nested structure.

Key Insight

A nested list behaves like a tree.

Lists are internal nodes.

Integers are leaves.

The depth increases by 1 whenever we go inside another list.

So the recursive function should carry the current depth:

dfs(current_list, depth)

It returns the weighted sum inside current_list.

Algorithm

Define a helper function:

dfs(items, depth)

For every item in items:

1. If item.isInteger() is true:
   1. Add item.getInteger() * depth.
2. Otherwise:
   1. Add dfs(item.getList(), depth + 1).

Return the total.

The final answer is:

dfs(nestedList, 1)

Walkthrough

Use:

nestedList = [1,[4,[6]]]

Start at depth 1.

The first element is 1.

It is an integer, so its contribution is:

1 * 1 = 1

The next element is [4,[6]].

It is a list, so we enter it at depth 2.

Inside that list, 4 contributes:

4 * 2 = 8

The next element is [6].

Enter it at depth 3.

Inside it, 6 contributes:

6 * 3 = 18

Total:

1 + 8 + 18 = 27

Correctness

The algorithm processes every element in the nested structure.

When it sees an integer at depth d, it adds exactly:

integer_value * d

which is exactly the contribution required by the problem.

When it sees a nested list at depth d, every integer inside that nested list is one level deeper. The recursive call uses d + 1, so all integers inside the nested list receive the correct depth.

The base case happens naturally when a list has no more elements to process. Its contribution is 0.

Since every integer belongs to exactly one position in the nested structure, the traversal counts each integer once. Therefore, the returned total is exactly the required weighted sum.

Complexity

Let N be the total number of NestedInteger objects, including integers and lists.

The maximum depth is bounded by the problem constraints.

Implementation

# """
# This is the interface that allows for creating nested lists.
# You should not implement it, or speculate about its implementation.
# """
# class NestedInteger:
#     def isInteger(self) -> bool:
#         ...
#
#     def getInteger(self) -> int:
#         ...
#
#     def getList(self) -> list["NestedInteger"]:
#         ...

class Solution:
    def depthSum(self, nestedList: list[NestedInteger]) -> int:
        def dfs(items: list[NestedInteger], depth: int) -> int:
            total = 0

            for item in items:
                if item.isInteger():
                    total += item.getInteger() * depth
                else:
                    total += dfs(item.getList(), depth + 1)

            return total

        return dfs(nestedList, 1)

Code Explanation

The helper receives a list and the depth of that list:

def dfs(items: list[NestedInteger], depth: int) -> int:

For each item, we check whether it is an integer:

if item.isInteger():

If yes, we add its weighted contribution:

total += item.getInteger() * depth

If the item is another list, we recurse one level deeper:

total += dfs(item.getList(), depth + 1)

The outermost list starts at depth 1:

return dfs(nestedList, 1)

Testing

LeetCode provides the NestedInteger interface, so local tests need a small mock class.

class NestedInteger:
    def __init__(self, value=None):
        self.value = value
        self.items = None if isinstance(value, int) else []

    def isInteger(self):
        return self.items is None

    def getInteger(self):
        return self.value

    def getList(self):
        return self.items

    def add(self, elem):
        if self.items is None:
            self.items = []
            self.value = None
        self.items.append(elem)

def build(value):
    if isinstance(value, int):
        return NestedInteger(value)

    node = NestedInteger()
    for item in value:
        node.add(build(item))
    return node

def run_tests():
    s = Solution()

    assert s.depthSum(build([[1, 1], 2, [1, 1]]).getList()) == 10
    assert s.depthSum(build([1, [4, [6]]]).getList()) == 27
    assert s.depthSum(build([0]).getList()) == 0
    assert s.depthSum(build([-1, [2, [-3]]]).getList()) == -6
    assert s.depthSum(build([[], [1]]).getList()) == 2

    print("all tests passed")

run_tests()

Test meaning: