Pythonのネストされたリストにおける重み付き合計 II(ボトムアップの重み付け)

問題概要

ネストされた整数のリストが与えられ、すべての整数を深さに基づく重みで掛けた合計を返すことを考えます。各要素は整数、またはリストであり、リストの要素にも整数や別のリストが含まれ得ます。

前回の問題では重みがルートから葉に向かって増加していましたが、今回は逆に下から上へ(ボトムアップ)重みが定義されます。つまり、最も深いレベル(葉)の整数の重みは 1 となり、ルートに近いレベルの整数ほど大きな重みを持ちます。

たとえば、入力が [[1,1],2,[1,1]] の場合、出力は 8 になります。これは、最深部にある 4 つの「1」が重み 1、その外側にある「2」が重み 2 となるためです(4×1 + 2×2 = 8)。

解法のアプローチ

この問題は、次の手順で解くことができます。

1. 関数 depthSumInverse() を定義し、ネストされたリスト nestedList を受け取ります。

2. flats:値と深さのペアを格納する新しいリストを用意します。

3. maxd:最大の深さを記録する変数を 0 で初期化します。

4. 補助関数 flatten() を定義します。引数はリスト nlst と現在の深さ dist です。

5. dist を 1 増やし、maxd を max(maxd, dist) で更新します。

6. nlst 内の各ノードについて処理を行います。

   • ノードが整数であれば、ペア (node, dist) を flats の末尾に追加します。
   • そうでなければ(リストであれば)、flatten(node, dist) を再帰的に呼び出します。

7. flatten(nestedList, 0) を呼び出し、リスト全体を「値と深さ」のペアにフラット化します。

8. summ を 0 で初期化します。

9. flats 内の各ペア (v, d) に対して、summ += v * (maxd + 1 - d) を計算します。

10. summ を返します。

ここで重要なのは、重みを maxd + 1 - d として計算する点です。これにより、最も深い要素(d = maxd)の重みがちょうど 1 になり、浅い階層の要素ほど大きな重みが与えられます。

実装例

理解を深めるために、以下の実装を見てみましょう。

class Solution(object):
    def depthSumInverse(self, nestedList):
        flats=[]
        self.maxd=0
        def flatten(nlst,dist):
            if isinstance(nlst,list):
                nlst=nlst
            dist+=1
            self.maxd=max(self.maxd,dist)
            for node in nlst:
                if isinstance(node,int):
                    flats.append((node,dist))
                else:
                    flatten(node,dist)
        flatten(nestedList,0)
        summ=0
        for v,d in flats:
            summ+=v*(self.maxd+1-d)
        return summ

ob = Solution()
print(ob.depthSumInverse([[1,1],2,[1,1]]))

入力

[[1,1],2,[1,1]]

出力

8

計算の流れ

入力 [[1,1],2,[1,1]] の場合、処理は次のように進みます。

• 外側のリストの深さは 1、内側のリストの深さは 2 となるため、maxd = 2 です。
• 4 つの「1」は深さ 2 にあるため、重みは 2 + 1 - 2 = 1。寄与は 4 × 1 × 1 = 4。
• 「2」は深さ 1 にあるため、重みは 2 + 1 - 1 = 2。寄与は 2 × 2 = 4。
• 合計は 4 + 4 = 8 となります。