floodflow.cc

#include <stack>
 
struct Index {
    int row; int col;
};
 
class Solution {
public:
    std::vector<std::vector<int>> floodFill(
        std::vector<std::vector<int>>& image, int sr, int sc, int color
    )
    {
        const int ORIGINAL_COLOR = image[sr][sc];
 
        if (ORIGINAL_COLOR == color) return image;
 
        const int ROWS = image.size();
        const int COLS = image[0].size();
 
        std::stack<Index> stk;
 
        Index source = {sr, sc};
        stk.push(source);
 
        while (!stk.empty()) {
            Index nodeIdx = stk.top();
            stk.pop();
 
            image[nodeIdx.row][nodeIdx.col] = color;
 
            if (nodeIdx.row + 1 < ROWS) {
                if (image[nodeIdx.row + 1][nodeIdx.col] == ORIGINAL_COLOR) {
                    Index nxt = { nodeIdx.row + 1, nodeIdx.col };
                    stk.push(nxt);
                }
            }
 
            if (nodeIdx.col + 1 < COLS) {
                if (image[nodeIdx.row][nodeIdx.col + 1] == ORIGINAL_COLOR) {
                    Index nxt = { nodeIdx.row, nodeIdx.col + 1 };
                    stk.push(nxt);
                }
            }
 
            if (nodeIdx.row - 1 >= 0) {
                if (image[nodeIdx.row - 1][nodeIdx.col] == ORIGINAL_COLOR) {
                    Index nxt = { nodeIdx.row - 1, nodeIdx.col };
                    stk.push(nxt);
                }
            }
 
            if (nodeIdx.col - 1 >= 0) {
                if (image[nodeIdx.row][nodeIdx.col - 1] == ORIGINAL_COLOR) {
                    Index nxt = { nodeIdx.row, nodeIdx.col - 1 };
                    stk.push(nxt);
                }
            }
 
        }
 
        return image;
    }
};