using optional

LifeGame.swift

extension Array {
    public subscript(safe index: Int) -> Element? {
        return indices.contains(index) ? self[index] : nil
    }
    
}

let gen: [[Int]] = [[1,1,0,0,0],
                    [1,1,0,0,0],
                    [0,1,0,0,0],
                    [0,0,0,0,0],
                    [0,0,0,0,0]]

func lifeGame(input generation: [[Int]]) -> [[Int]] {
    
    let lineArray = Array.init(repeating: 0, count: generation.first?.count ?? 0)
    var nextGeneration: [[Int]] = Array.init(repeating: lineArray, count: generation.count)
    var count: Int = 0
    
    for x in 0..<generation.count {
        for y in 0..<generation.count {x
            
            // above
            if let cell = generation[safe: x-1]?[safe: y-1], cell == 1 { count += 1 }
            if let cell = generation[safe: x-1]?[safe: y], cell == 1 { count += 1 }
            if let cell = generation[safe: x-1]?[safe: y+1], cell == 1 { count += 1 }
            // left right
            if let cell = generation[safe: x]?[safe: y-1], cell == 1 { count += 1 }
            if let cell = generation[safe: x]?[safe: y+1], cell == 1 { count += 1 }
            // under
            if let cell = generation[safe: x+1]?[safe: y-1], cell == 1 { count += 1 }
            if let cell = generation[safe: x+1]?[safe: y], cell == 1 { count += 1 }
            if let cell = generation[safe: x+1]?[safe: y+1], cell == 1 { count += 1 }
            
            // check to survive
            if (generation[x][y] == 0 && count == 3) || (generation[x][y] == 1 && (count == 2 || count == 3)){
                nextGeneration[x][y] = 1
            }
            
            count = 0 // reset count
        }
    }
    
    return nextGeneration
}