C++で解く0と1のマトリックス問題：BFSで各セルから最も近い0までの距離を求める方法

問題の概要

0と1だけで構成された行列（マトリックス）が与えられたとき、各セルについて「最も近い0」までの距離を求める問題を考えます。ここで、隣接する2つのセル間の距離は1と定義します。

例として、次のような入力が与えられたとします。

この場合、出力は次のようになります。

中央の「1」は上下左右のいずれかに存在する0までの距離が1であり、左下の「1」とその右隣の「1」は、それぞれ最寄りの0までの距離が1と2であることを示しています。

解法のアプローチ

この問題は幅優先探索（BFS）を用いることで効率的に解くことができます。基本的な考え方は、まずすべての「0」のセルを探索の起点としてキューに登録し、そこから距離1、距離2…と外側へ層状（レベルごと）に探索範囲を広げていくというものです。すでにより近い0が見つかっているセルには更新を行わないことで、無駄な計算を避けられます。

アルゴリズムの手順

• 移動方向を表すサイズ4×2の配列 dir := {{1, 0}, {-1, 0}, {0, -1}, {0, 1}} を定義します。
• n := 行数、m := 列数とします。
• n × m の結果格納用行列 ret を定義し、すべて無限大（inf）で初期化します。
• キュー q を用意します。
• i := 0 から n 未満の間、i を1ずつ増やしながら以下を繰り返します。
  • j := 0 から m 未満の間、j を1ずつ増やしながら以下を繰り返します。
    • matrix[i][j] が 0 の場合：
      • ret[i][j] := 0 と設定する
      • 座標 {i, j} をキュー q に挿入する
• lvl := 1 とし、キュー q が空でない限り、lvl を1ずつ増やしながら以下を繰り返します。
  • sz := キュー q の現在のサイズとします。
  • sz が 0 になるまで（各反復で sz を1減らしながら）以下を繰り返します。
    • curr := キューの先頭要素（座標ペア）を取得し、キューから取り除きます。
    • k := 0 から 4 未満の間、k を1ずつ増やしながら以下を繰り返します。
      • nx := curr.first + dir[k][0]
      • ny := curr.second + dir[k][1]
      • nx < 0、nx >= n、ny < 0、ny >= m のいずれかを満たす場合、またはすでに ret[nx][ny] < lvl で確定済みの場合はスキップします。
      • それ以外の場合は ret[nx][ny] := lvl と設定し、{nx, ny} をキュー q に挿入します。
• 最後に ret を返します。

C++による実装例

より理解を深めるために、以下の実装例をご覧ください。

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<vector<auto> > v){
    cout << "[";
    for(int i = 0; i<v.size(); i++){
       cout << "[";
       for(int j = 0; j <v[i].size(); j++){
          cout << v[i][j] << ", ";
       }
       cout << "],";
    }
    cout << "]"<<endl;
}
int dir[4][2] = {{1, 0}, {-1, 0}, {0, -1}, {0, 1}};
class Solution {
public:
    vector<vector<int>> updateMatrix(vector<vector<int>>& matrix) {
    int n = matrix.size();
    int m = matrix[0].size();
    vector < vector <int> > ret(n, vector <int>(m, INT_MAX));
    queue < pair <int, int> > q;
    for(int i = 0; i < n; i++){
       for(int j = 0; j < m; j++){
          if(!matrix[i][j]){
             ret[i][j] = 0;
             q.push({i, j});
          }
       }
    }
    for(int lvl = 1; !q.empty(); lvl++){
       int sz = q.size();
       while(sz--){
          pair <int, int> curr = q.front();
          q.pop();
          for(int k = 0; k < 4; k++){
             int nx = curr.first + dir[k][0];
             int ny = curr.second + dir[k][1];
             if(nx < 0 || nx >= n || ny < 0 || ny >= m || ret[nx][ny] < lvl) continue;
                ret[nx][ny] = lvl;
                q.push({nx, ny});
            }
          }
       }
       return ret;
    }
};
main(){
    Solution ob;
    vector<vector<int>> v = {{0,0,0},{0,1,0},{1,1,1}};
    print_vector(ob.updateMatrix(v));
}

入力

{{0,0,0},{0,1,0},{1,1,1}}

出力

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

まとめ

このように、BFSを活用することで、各セルから最も近い0までの距離を効率的に計算できます。計算量は行数・列数をそれぞれn、mとするとO(n × m)となり、大きな行列に対しても高速に動作するのが特徴です。類似のグリッド探索問題にも応用できるテクニックなので、ぜひマスターしておきましょう。