Pythonで捕まえられる雨水の総量を計算するプログラム（トレッピング・レイン・ウォーター問題）

非負整数からなる長さ n の配列が与えられているとします。各要素はバーの高さを表し、それぞれのバーの幅は1です。このとき、雨が降った後に溜め込むことのできる水の総量を計算するのが本記事のテーマです。状況を図にすると、以下のようになります。

図を見ると、水が溜まっている部分（青い箱）は全部で8個あります。したがって、このケースの出力は8となります。

解法のアプローチ

この問題は「スタック」を利用することで効率よく解けます。全体の手順は以下の通りです。

1. スタック st、変数 water := 0、インデックス i := 0 を用意する
2. i が高さ配列のサイズ未満である間、次の処理を繰り返す
   • スタックが空であるか、height[スタックの先頭] >= height[i] ならば、i をスタックにプッシュし、i を1増やす
   • それ以外の場合は次のように処理する
     • x := スタックの先頭要素を取得し、スタックからポップする
     • スタックが空でない場合は以下を計算する
       • temp := height[スタックの先頭] と height[i] のうち小さい方
       • dist := i − スタックの先頭 − 1
       • water := water + dist × (temp − height[x])
3. 最後に water を返す

このアルゴリズムでは、スタックにインデックスを記録しながら走査することで、左右どちら側にも高い壁がある凹み部分に溜まる水を段階的に積算していきます。計算量は O(n)、空間計算量も O(n) であり、非常に効率的な手法として知られています。

それでは、実際のPythonによる実装例を見て理解を深めましょう。

実装例

class Solution(object):
   def trap(self, height):
      stack = []
      water = 0
      i=0
      while i<len(height):
         if len(stack) == 0 or height[stack[-1]]>=height[i]:
            stack.append(i)
            i+=1
         else:
            x = stack[-1]
            stack.pop()
            if len(stack) != 0:
               temp = min(height[stack[-1]],height[i])
               dist = i - stack[-1]-1
               water += dist*(temp - height[x])
      return water
ob = Solution()
print(ob.trap([2,5,2,0,5,8,8]))

入力

[2,5,2,0,5,8,8]

出力

8

このように、スタックを使ったアプローチを採用すれば、与えられた高さの配列に対して溜まる雨水の総量を正確に求めることができます。ぜひ自分のコードでも試してみてください。