nested_list_weight_sum

 1from collections import defaultdict
 2
 3
 4class NestedInteger:
 5    """@private"""
 6
 7    def __init__(self, value=None):
 8        self.value = []
 9        if value:
10            self.value = value
11
12    def isInteger(self) -> bool:
13        return self.value is int
14
15    def add(self, elem: int):
16        self.value.append(elem)
17
18    def setInteger(self, value):
19        self.value = value
20
21    def getInteger(self):
22        if self.value is int:
23            return self.value
24        raise RuntimeError("Get List must be called on a NestedInteger that is an int")
25
26    def getList(self) -> list:
27        if self.value is list:
28            return self.value
29        raise RuntimeError("Get List must be called on a NestedInteger that is a list")
30
31
32# @leet start
33class Solution:
34    def depthSum(self, nestedList: list[NestedInteger]) -> int:
35        """
36        This problem asks to implement a function that implements the sum for a list of `NestedInteger`s.
37        The NestedInteger class can either be an integer or a list.
38        The total of a List of NestedInteger is each NestedInteger multiplied by its depth, summed up.
39        For example, take this list: `[[1, 1], 2, [1, 1]]`. There are four integers with depth 2, which are 1
40        and one 2, which has depth 1. So, the total is 10.
41
42        To solve this, we can first iterate through the list to determine each item's depth, and assign it to a dict.
43        The depth has to be found either iteratively or recursively. This solution does it recursively.
44        For each item, the function calls a subroutine that checks if the current number is an integer.
45        If it is, it adds it to the dictionary with its current depth.
46
47        If it is not, then for each item in the sublist, it recurses, adding 1 to its depth.
48
49        Finally, each key -> value pair of depth -> numbers is multiplied together, and summed up.
50        """
51        depths = defaultdict(list)
52
53        def recurse(curr: NestedInteger, depth: int):
54            if curr.isInteger():
55                depths[depth].append(curr.getInteger())
56            else:
57                for item in curr.getList():
58                    recurse(item, depth + 1)
59
60        for item in nestedList:
61            recurse(item, 1)
62
63        return sum(k * sum(v) for k, v in depths.items())
64
65
66# @leet end
67def test():
68    pass