func countries(in grid: inout [[Int]]) -> Int {
    let height = grid.count
    let width = grid[0].count

    var flag = 0
    var flagCount = 0

    func cascade(x: Int, y: Int) {
        guard grid[y][x] == flag else {
            return
        }

        grid[y][x] = 0

        if x > 0 {
            cascade(x: x - 1, y: y)
        }
        if x < width - 1 {
            cascade(x: x + 1, y: y)
        }
        if y > 0 {
            cascade(x: x, y: y - 1)
        }
        if y < height - 1 {
            cascade(x: x, y: y + 1)
        }
    }

    for x in 0..<width {
        nextFlag: for y in 0..<height {
            flag = grid[y][x]

            if flag == 0 { continue nextFlag }

            flagCount += 1

            cascade(x: x, y: y)
        }
    }

    return flagCount
}

var grid: [[Int]] = [
    [1, 1, 1, 3, 1],
    [1, 2, 2, 3, 1],
    [2, 1, 3, 3, 1],
    [1, 1, 4, 4, 4]
]

let count = countries(in: &grid)
print(count)
assert(count == 7)