[LintCode] Problem 551 - Nested List Weight Sum

Posted on | In ,

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.

Example

No.1

Input: the list [[1,1],2,[1,1]],

Output: 10.

Explanation:
four 1’s at depth 2, one 2 at depth 1, 4 * 1 * 2 + 1 * 2 * 1 = 10

No.2

Input: the list [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

Code

1
2
3
4
5
public interface NestedInteger {
    public boolean isInteger();
    public Integer getInteger();
    public List<NestedInteger> getList();
}
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
public int depthSum(List<NestedInteger> nestedList) {
    Queue<NestedInteger> queue = new LinkedList<>();
    int level = 0;
    int sum = 0;

    for (NestedInteger nestedInteger : nestedList)
        queue.offer(nestedInteger);

    while (!queue.isEmpty()) {
        int size = queue.size();
        level++;

        for (int i = 0; i < size; i++) {
            NestedInteger nestedInteger = queue.poll();

            if (nestedInteger.isInteger())
                sum += level * nestedInteger.getInteger();
            else {
                for (NestedInteger list : nestedInteger.getList())
                    queue.offer(list);
            }
        }
    }

    return sum;
}