Swift で Flood Fill（塗りつぶし）

備忘用メモ

NOTE

• 実装：dfs
• 問題
  • https://leetcode.com/problems/flood-fill/
  • グリッド内で任意のセルを開始地点にして４方向に隣接したセルを塗りつぶす

コード

// sr: row  of graph
// sc: cell of graph
func floodFill(_ image: [[Int]], _ sr: Int, _ sc: Int, _ newColor: Int) -> [[Int]] {
    // ready
    let directions = [
        [ 0, -1], // top
        [ 1,  0], // right
        [ 0,  1], // bottom
        [-1,  0], // left
    ]
    let lenX = image.count
    let lenY = image[0].count
    var visited = [[Bool]](repeating: [Bool](repeating: false, count: lenY), count: lenX)
    let prevColor = image[sr][sc]

    // define dfs
    func dfs(_ graph: inout [[Int]], _ sr: Int, _ sc: Int) {
        if sr < 0 || sr > lenX-1 || sc < 0 || sc > lenY-1 { return } // invalid param
        if graph[sr][sc] != prevColor { return }
        if visited[sr][sc] { return }
        visited[sr][sc] = true
        graph[sr][sc] = newColor
        for dir in directions {
            let x = dir[0], y = dir[1]
            dfs(&graph, sr + x, sc + y)
        }
    }

    // run
    var _image = image
    dfs(&_image, sr, sc)
    return _image
}

ポイント

• 隣接セルへの移動は directions 変数で管理

  • 移動先を変数として定義しておく。この問題では隣接セルが４セルのことを指していたが、たまに８セルのケースもあるのでその場合は directions 変数を変更すればよいだけになる

• 無限ループ回避のために visited 変数で訪問済みかチェック（この問題の場合はなくても問題ないが基本的には必要）