An iterative approach to solving a variation of the flood fill problem in swift

Contiguous Colors In Grid: Iterative.swift

extension Array {
    public subscript(safe index: Int) -> Element? {
        guard index >= 0, index < endIndex else {
            return nil
        }

        return self[index]
    }
}

/*

 Given a unknown size matrix of colors, return the largest contiguous number of a single color

 🔵  🟢  🔴

 🟢  🔴  🔴

 🔴  🔴  🔴

 => 6 total red points

 */

final class Point {

    enum Color {
        case red
        case green
        case blue
    }

    let x, y: Int
    let color: Color
    var hasVisited: Bool

    init(x: Int, y: Int, color: Color, hasVisited: Bool = false) {
        self.x = x
        self.y = y
        self.color = color
        self.hasVisited = hasVisited
    }

    func top(in grid: Grid) -> Point? {
        grid.points[safe: y + 1]?[x]
    }

    func left(in grid: Grid) -> Point? {
        grid.points[safe: y]?[safe: x - 1]
    }

    func right(in grid: Grid) -> Point? {
        grid.points[safe: y]?[safe: x + 1]
    }

    func bottom(in grid: Grid) -> Point? {
        grid.points[safe: y - 1]?[safe: x]
    }

    func allValidNeighbors(in grid: Grid) -> [Point] {
        let neighbors = [
            top(in: grid),
            left(in: grid),
            right(in: grid),
            bottom(in: grid)
        ]
        let valid = neighbors.reduce(into: [Point]()) { neighbors, neighbor in
            guard let neighbor = neighbor, neighbor.color == color else { return }
            neighbors.append(neighbor)
        }

        return valid
    }
}

struct Grid {

    let points: [[Point]]

    func resetToDefault() {
        points.flatMap({ $0 }).forEach {
            $0.hasVisited = false
        }
    }

    // MARK: - Iterative
    func largestContiguousPointSet() -> Int {

        var stack = points.flatMap { $0 }
        var colorCount = [Point.Color:Int]()

        repeat {
            let point = stack.removeLast()

            for neighbor in point.allValidNeighbors(in: self) where neighbor.hasVisited == false {
                colorCount[neighbor.color] = (colorCount[neighbor.color] ?? 0) + 1
                neighbor.hasVisited = true
                stack.append(contentsOf: neighbor.allValidNeighbors(in: self))
            }

        } while stack.isEmpty == false

        return colorCount.values.max() ?? 0
    }

    // MARK: - Recursive
    func largestContiguiousPointSetRecursive() -> Int {
        return contiguousArea(for: points.flatMap{ $0 }).values.max() ?? 0
    }

    private func contiguousArea(for points: [Point]) -> [Point.Color:Int] {
        return points.reduce(into: [Point.Color:Int]()) { result, point in
            guard point.hasVisited == false else { return }
            result[point.color] = (result[point.color] ?? 0) + 1
            point.hasVisited = true
            for (key, value) in contiguousArea(for: point.allValidNeighbors(in: self)) {
                result[key] = (result[key] ?? 0) + value
            }
        }
    }
}

let grid = Grid(points: [
    [
        .init(x: 0, y: 0, color: .red),
        .init(x: 1, y: 0, color: .red),
        .init(x: 2, y: 0, color: .red),
    ],
    [
        .init(x: 0, y: 1, color: .green),
        .init(x: 1, y: 1, color: .red),
        .init(x: 2, y: 1, color: .red),
    ],
    [
        .init(x: 0, y: 2, color: .blue),
        .init(x: 1, y: 2, color: .green),
        .init(x: 2, y: 2, color: .red),
    ]
])

grid.largestContiguousPointSet() // => 6
grid.resetToDefault()
grid.largestContiguiousPointSetRecursive() // => 6