Skip to content

339. Nested List Weight Sum

Approach 1: BFS

  • Time: $O(n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class Solution {
 public:
  int depthSum(vector<NestedInteger>& nestedList) {
    int ans = 0;
    int depth = 0;
    queue<NestedInteger> q;

    addIntegers(q, nestedList);

    while (!q.empty()) {
      ++depth;
      for (int sz = q.size(); sz > 0; --sz) {
        const NestedInteger ni = q.front();
        q.pop();
        if (ni.isInteger())
          ans += ni.getInteger() * depth;
        else
          addIntegers(q, ni.getList());
      }
    }

    return ans;
  }

 private:
  void addIntegers(queue<NestedInteger>& q,
                   const vector<NestedInteger>& nestedList) {
    for (const NestedInteger& ni : nestedList)
      q.push(ni);
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class Solution {
  public int depthSum(List<NestedInteger> nestedList) {
    int ans = 0;
    int depth = 0;
    Queue<NestedInteger> q = new ArrayDeque<>();

    addIntegers(q, nestedList);

    while (!q.isEmpty()) {
      ++depth;
      for (int sz = q.size(); sz > 0; --sz) {
        final NestedInteger ni = q.poll();
        if (ni.isInteger())
          ans += ni.getInteger() * depth;
        else
          addIntegers(q, ni.getList());
      }
    }

    return ans;
  }

  private void addIntegers(Queue<NestedInteger> q, List<NestedInteger> nestedList) {
    for (final NestedInteger ni : nestedList)
      q.offer(ni);
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution:
  def depthSum(self, nestedList: list[NestedInteger]) -> int:
    ans = 0
    depth = 0
    q = collections.deque()

    def addIntegers(nestedList: list[NestedInteger]) -> None:
      for ni in nestedList:
        q.append(ni)

    addIntegers(nestedList)

    while q:
      depth += 1
      for _ in range(len(q)):
        ni = q.popleft()
        if ni.isInteger():
          ans += ni.getInteger() * depth
        else:
          addIntegers(ni.getList())

    return ans

Approach 2: DFS

  • Time: $O(n)$
  • Space: $O(h)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
class Solution {
 public:
  int depthSum(vector<NestedInteger>& nestedList) {
    int ans = 0;
    dfs(nestedList, 1, ans);
    return ans;
  }

 private:
  void dfs(const vector<NestedInteger>& nestedList, int depth, int& ans) {
    for (const NestedInteger& ni : nestedList)
      if (ni.isInteger())
        ans += ni.getInteger() * depth;
      else
        dfs(ni.getList(), depth + 1, ans);
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
class Solution {
  public int depthSum(List<NestedInteger> nestedList) {
    dfs(nestedList, 1);
    return ans;
  }

  private int ans = 0;

  private void dfs(final List<NestedInteger> nestedList, int depth) {
    for (final NestedInteger ni : nestedList)
      if (ni.isInteger())
        ans += ni.getInteger() * depth;
      else
        dfs(ni.getList(), depth + 1);
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
class Solution:
  def depthSum(self, nestedList: list[NestedInteger]) -> int:
    ans = 0

    def dfs(nestedList: list[NestedInteger], depth: int) -> None:
      nonlocal ans
      for ni in nestedList:
        if ni.isInteger():
          ans += ni.getInteger() * depth
        else:
          dfs(ni.getList(), depth + 1)

    dfs(nestedList, 1)
    return ans