C++で2次元地形の谷間に溜まる雨水の量を求めるプログラム

2次元行列があり、各要素は地形の高さを表しているものとします。ここで雨が降り、谷間のくぼみがすべて水で満たされる状況を想像してみてください。

私たちの課題は、谷間に溜まる雨水の総量を求めることです。

たとえば、入力が次のような場合を考えてみましょう。

この場合の出力は 3 になります。中央の高さ4のマスには2単位、その隣の高さ5のマスには1単位の水が溜まり、合計3単位となるためです。

アルゴリズムのアプローチ

この問題は「外周から内側へ」処理を進めることで効率的に解けます。水は必ず外側へあふれるため、外周を起点に水位(せきの高さ)を追高さ)を追跡しながら内部へ進んでいくイメージです。具体的な手順は以下の通りです。

• 座標x・yと高さhを持つ構造体Dataを定義します。
• 高さを基準に並べ替えられる優先度付きキューpq(最小ヒープとして動作)を用意します。
• n := hの行数とします。nが0の場合は0を返します。
• m := h[0]の列数とします。
• 訪問済みセルを記録する集合visitedを定義します。
• まず行列の外周(四辺)にあるセルをすべてpqに追加し、同時にvisitedにも登録します。
• ret := 0、maxVal := 0 で初期化します。
• pqが空になるまで、次の処理を繰り返します。
  • pqの先頭(最も高さの小さい)要素tempを取り出します。
  • maxValをtempの高さとの最大値で更新します。これが現在の「せき」の高さになります。
  • tempの座標をx、yとします。
  • iを0から3まで動かしながら、上下左右4方向の隣接セル(nx, ny)を調べます。
  • nx・nyが行列の範囲内で、かつ未訪問の場合:
    • val := h[nx][ny] とします。
    • valがmaxValより小さければ、その差(maxVal − val)だけ水が溜まるとみなし、retに加算したうえでvalをmaxValに引き上げます。
    • Data(val, nx, ny)をpqに追加し、{nx, ny}をvisitedに登録します。
• 最後にretを返します。

実装例

それでは、実際のC++コードを見ながら理解を深めましょう。

#include <bits/stdc++.h>
using namespace std;
struct Data {
    int x, y;
    int h;
    Data(int a, int b, int c) {
        h = a;
        x = b;
        y = c;
    }
};
struct Comparator {
    bool operator()(Data a, Data b) {
        return !(a.h < b.h);
    }
};
int dir[4][2] = {{1, 0}, {-1, 0}, {0, -1}, {0, 1}};
class Solution {
   public:
   int solve(vector<vector<int>>& h) {
      priority_queue<Data, vector<Data>, Comparator> pq;
      int n = h.size();
      if (!n)
         return 0;
      int m = h[0].size();
      set<pair<int, int>> visited;
      for (int i = 0; i < n; i++) {
         pq.push(Data(h[i][0], i, 0));
         visited.insert({i, 0});
         pq.push(Data(h[i][m - 1], i, m - 1));
         visited.insert({i, m - 1});
      }
      for (int i = 1; i < m - 1; i++) {
         pq.push(Data(h[0][i], 0, i));
         visited.insert({0, i});
         pq.push(Data(h[n - 1][i], n - 1, i));
         visited.insert({n - 1, i});
      }
      int ret = 0;
      int maxVal = 0;
      while (!pq.empty()) {
         Data temp = pq.top();
         pq.pop();
         maxVal = max(temp.h, maxVal);
         int x = temp.x;
         int y = temp.y;
         int nx, ny;
         for (int i = 0; i < 4; i++) {
            nx = x + dir[i][0];
            ny = y + dir[i][1];
            if (nx >= 0 && ny >= 0 && nx < n && ny < m && !visited.count({nx, ny})) {
               int val = h[nx][ny];
               if (val < maxVal) {
                  ret += maxVal - val;
                  val = maxVal;
               }
               pq.push(Data(val, nx, ny));
               visited.insert({nx, ny});
            }
         }
      }
      return ret;
   }
};
int solve(vector<vector<int>>& matrix) {
   return (new Solution())->solve(matrix);
}
int main(){
   vector<vector<int>> v = {
      {6, 6, 6, 8},
      {6, 4, 5, 8},
      {6, 6, 6, 6}
   };
   cout << solve(v);
}

入力

{
   {6, 6, 6, 8},
   {6, 4, 5, 8},
   {6, 6, 6, 6}
};

出力

3

計算量について

各セルは最大1回ずつキューに追加され、ヒープ操作1回あたりO(log(nm))のコストがかかるため、全体の時間計算量はO(nm log(nm))です。また、訪問済みセルの管理などに必要な空間計算量はO(nm)となります。この手法を使えば、大規模な地形データでも効率よく溜まる雨水の量を求められます。