Взвешенная сумма вложенного списка

Источник: https://leetcode.com/problems/nested-list-weight-sum-ii/

Дан вложенный список целых чисел. Подсчитать сумму всех его элементов, взвешенных их глубиной.

Каждый элемент вложенного списка может быть целым числом или вложенным списком - тоже состоящий из целых чисел и вложенных списков.

Вес увеличивается снизу вверх, то есть "листьевые" целые элементы имеют вес 1. А корень имеет наибольший вес.

Как изменится алгоритм, если вес увеличивается от корня к листьям?

Пример 1:

Дано: [[1,1],2,[1,1]]
Результат: 8 
Пояснение: Четыре единицы находятся на глубине 1, а 2 - на глубине 2. Поэтому сумма считается: 1*1 + 1*1 + 2*2 + 1*1 + 1*1 = 8.

Пример 2:

Дано: [1,[4,[6]]]
Результат: 17 
Поясненеи: Число 1 находится на глубине 3, 4 - на глубине 2, 6 - на глубине 1. Поэтому сумма считается: 1*3 + 4*2 + 6*1 = 17.

Есть некоторый класс NestedInteger, со следующим интерфейсом:

/**
 * // Интерфейс для вложенного списка, не нужно его имплементировать
 *
 * public interface NestedInteger {
 *     // Консструктор создает пустой список
 *     public NestedInteger();
 *
 *     // Constructor инициализирует список из одного целого числа.
 *     public NestedInteger(int value);
 *
 *     // @return true если рассамтриваемый элемент NestedInteger содержит целое число
 *     public boolean isInteger();
 *
 *     // @return целочисленное значение, если NestedInteger содержит целое число
 *     // возвращает null если элемент NestedInteger является вложенным списком
*     // (а не целым числом)
 *     public Integer getInteger();
 *
 *     // Элемент NestedInteger содержит целое число.
 *     public void setInteger(int value);
 *
 *     // Если NestedInteger - это список, то к нему добавляется целое число или список
 *     public void add(NestedInteger ni);
 *
 *     // @return вложенный список, который находится в элементе NestedInteger, если это список
 *     // Return null if this NestedInteger holds a single integer
 *     public List<NestedInteger> getList();
 * }
 */

Идея решения

Мы не будем подсчитывать глубину дерева, которое сформировано вложенным списком. 

Мы пробегаем по элементам вложенного списка. Если элемент список - мы добавляем его на следующую итерацию. Если целое число - мы суммируем его в сумму текущей итерации.

Сумма целых чисел - это невзвешенная (unweighted) сумма всех целых чисел, которые мы встретили на текущей итерации. Мы её не обнуляем ни на каком уровне. 

Например, на корневом уровне мы подсчитали sum1. На следующем уровне мы нашли  sum2 целых элементов. Затем sum3. И так далее.

Но если мы её не обнуляем, а аккумулируем, в виде:

sum1 + (sum2 + sum1) + (sum3 + sum2 + sum1) +...

то получится, что самые вперхние элементы, находящиеся в sum1, будут умножены на глубину дерева. Затем целые элементы, находящиеся на втором уровне, будут повторяться в sum2, d-1 раз, если d - это глубина дерева. И так далее. Самые нижние элементы будут подсчитаны один раз, то есть sumd войдет в финальную сумму 1 раз, то есть с весом 1.

Имплементация

class Solution {
    public int depthSumInverse(List<NestedInteger> nestedList) {
        int unweighted = 0, weighted = 0;
        while (!nestedList.isEmpty()) {
// если элемент не целое число, мы его отложим на следюущую итерацию
            List<NestedInteger> nextLevel = new ArrayList<>();
            for (NestedInteger ni : nestedList) {
                if (ni.isInteger()) // формируем sumx, x - очередной "этаж" дерева сверху вниз
                    unweighted += ni.getInteger();
                else
                    nextLevel.addAll(ni.getList());
            }
// в unweighted мы подсчитали  (sum3 + sum2 + sum1)
// в weighted уже хранится sum1 + (sum2 + sum1)
            weighted += unweighted;
            nestedList = nextLevel;
        }
        return weighted;
    }
}

Если вес увеличвается от корня к листьям, то алгоритм становится гораздо проще: на каждом уровне мы увеличиваем счётчик глубины. И не взвешенная сумма целых чисел просто умножается на глубину и добавляется к фмнальному результату.

public int depthSum(List<NestedInteger> nestedList) {
        int weighted = 0; // финальная взвешенная сумма, финальный результат
        int d=0; // счетчик глубины
        while (!nestedList.isEmpty()) {
            int unweighted=0; // невзвешенная сумма обнуляется на каждом уровне
            List<NestedInteger> nextLevel = new ArrayList<>();
            for (NestedInteger ni : nestedList) {
                if (ni.isInteger())
                    unweighted += ni.getInteger();
                else
                    nextLevel.addAll(ni.getList());
            }
            d++;
            weighted += unweighted * d; // невзвешенная сумма умножается на глубину
                        // и суммируется с финальным результатом
            nestedList = nextLevel;
        }
        return weighted;
    }