C++で電球が照らせるセルの最大数を求めるプログラム

問題の概要

高さ h × 幅 w のグリッドが与えられます。グリッドの各セルには電球または障害物が置かれており、電球のあるセルは上下左右のセルを照らします。光は障害物に遮られない限りまっすぐ進みますが、障害物のセル自体は照らすことができず、電球の光が他のセルへ届くのも遮ります。グリッドは文字列の配列として渡され、「#」が障害物、「.」が空きセルを表します。手持ちの電球は1つだけなので、グリッド内の最適な位置に配置したときに照らせるセルの最大数を求める必要があります。

たとえば、入力が h = 4、w = 4、grid = {"#...", "....", "...#", "...."} の場合、出力は 7 になります。

上の図から、グリッド内で実際に照らされているセルを確認できます。

解き方の手順

この問題を解くために、次の手順に従います。

2次元配列 first を定義する
i := 0 で初期化し、i < h の間、(i を1ずつ増やしながら) 以下を繰り返す:
    count := 0
    j := 0 で初期化し、j < w の間、(j を1ずつ増やしながら) 以下を繰り返す:
        grid[i, j] が '#' と等しい場合:
            count := 0
            以降の処理をスキップして次の反復へ進む
        それ以外の場合:
            first[i, j] := count
            (count を1増やす)
    k := 0
    j := w - 1 で初期化し、j >= 0 の間、(j を1ずつ減らしながら) 以下を繰り返す:
        grid[i, j] が '#' と等しい場合:
            k := 0
            以降の処理をスキップして次の反復へ進む
        それ以外の場合:
            k := k と first[i, j] の最大値
            first[i, j] := k
2次元配列 second を定義する
j := 0 で初期化し、j < w の間、(j を1ずつ増やしながら) 以下を繰り返す:
    count := 0
    i := 0 で初期化し、i < h の間、(i を1ずつ増やしながら) 以下を繰り返す:
        grid[i, j] が '#' と等しい場合:
            count := 0
            以降の処理をスキップして次の反復へ進む
        それ以外の場合:
            second[i, j] := count
            (count を1増やす)
    k := 0
    i := h - 1 で初期化し、i >= 0 の間、(i を1ずつ減らしながら) 以下を繰り返す:
        grid[i, j] が '#' と等しい場合:
            k := 0
            以降の処理をスキップして次の反復へ進む
        それ以外の場合:
            k := k と second[i, j] の最大値
            second[i, j] := k
result := 0
i := 0 で初期化し、i < h の間、(i を1ずつ増やしながら) 以下を繰り返す:
    j := 0 で初期化し、j < w の間、(j を1ずつ増やしながら) 以下を繰り返す:
        result := result と first[i, j] + second[i, j] の最大値
result + 1 を返す

アルゴリズムの考え方

まず、2次元配列 first を用意し、各行を左から右へ走査しながら、障害物にぶつかるまで連続する空きセルの数を数えます。続けて同じ行を右から左へ走査し、それまでの最大値を各セルに書き戻すことで、そのセルに電球を置いた場合に左右方向へ照らせるセルの総数が first[i][j] に格納されます。

同様に、2次元配列 second では列ごとに上下方向へ照らせるセル数を計算します。

最後に、すべてのセルについて first[i][j] + second[i][j] を調べ、その最大値に電球自身のセルを表す 1 を加えたものが答えとなります。この方法なら、各セルを定数回走査するだけでよく、計算量は O(h × w) と非常に効率的です。

実装例

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

#include <bits/stdc++.h>
using namespace std;

int solve(int h, int w, vector<string> grid){
    vector<vector<int>> first(h, vector<int> (w));
    for(int i = 0; i < h; i++) {
        int count = 0;
        for(int j = 0; j < w; j++) {
            if(grid[i][j] == '#') {
                count = 0;
                continue;
            } else {
                first[i][j] = count;
                count++;
            }
        }
        int k = 0;
        for(int j = w-1; j >= 0; j--) {
            if(grid[i][j] == '#') {
                k = 0;
                continue;
            } else {
                k = max(k, first[i][j]);
                first[i][j] = k;
            }
        }
    }
    vector<vector<int>> second(h, vector<int> (w));
    for(int j = 0; j < w; j++) {
        int count = 0;
        for(int i = 0; i < h; i++) {
            if(grid[i][j] == '#') {
                count = 0;
                continue;
            } else {
                second[i][j] = count;
                count++;
            }
        }
        int k = 0;
        for(int i = h-1; i >= 0; i--) {
            if(grid[i][j] == '#') {
                k = 0;
                continue;
            } else {
                k = max(k, second[i][j]);
                second[i][j] = k;
            }
        }
    }
    int result = 0;
    for(int i = 0; i < h; i++) {
        for(int j = 0; j < w; j++) {
            result = max(result, first[i][j] + second[i][j]);
        }
    }
    return result + 1;
}
int main() {
   int h = 4, w = 4;
   vector<string> grid = {"#...", "....", "...#", "...."};
   cout<< solve(h, w, grid);
   return 0;
}

入力

4, 4, {"#...", "....", "...#", "...."}

出力

7