This snippet finds the largest square area within a matrix with consisting of 1s.

LargestSquareArea.swift

import Foundation
import XCTest

typealias Matrix = [[Int]]

/// This function alias is used when a position is given:
/// And we'd like to check if we can go to a neighbor of it
/// And if we can go to the neighbor, what's the value of the neighbor
/// An example is usage:
/// ```
/// let matrix = [
///   [0, 0, 0],
///   [1, 1, 1],
///   [1, 1, 0
/// ]
/// let position = Position(row:0, col:0)
/// let (canGo, valueOfRightNeighbor) = position.hasRightNeighbor(in:matrix)
/// print(canGo) // prints true
/// print(valueOfRightNeighbor) // prints 0
///
/// let allMovements = position.allMovementsPossible // returns bunch of movements, top, left etc.
/// for move in allMovements {
///     print(move(matrix)) => prints (Bool, Int?) result
/// }
/// ```
typealias CanMoveValueOfMoveFromPositionFunctionReturnType = (canMove: Bool, valueOfNeighbor:Int?)
typealias CanMoveValueOfMoveFromPositionFunctionType =
    ((Matrix) -> CanMoveValueOfMoveFromPositionFunctionReturnType)

// MARK: - ProductType
/// Represents product types
enum ProductType: Int {
    case nonDefect = 0
    case defect = 1
}

// MARK: - Position
/// Represents a position in two dimensional matrix
struct Position: Equatable, Comparable {
    // MARK: - Properties
    let row: Int
    let col: Int
    
    // MARK: - Equatable
    static func == (lhs: Position, rhs: Position) -> Bool {
        return lhs.row == rhs.row && lhs.col == rhs.col
    }
    
    // MARK: - Comparable
    static func < (lhs: Position, rhs: Position) -> Bool {
        return lhs.row < rhs.row || lhs.col < rhs.col
    }
    
    // This solution uses diagonal right bottom advancement algorithm. Hence, only enabled moves
    // are right, bottom, and bottom right.
    func allMovementsPossible(in matrix:Matrix) -> [CanMoveValueOfMoveFromPositionFunctionType] {
        return [
            hasRightNeighbor(in:),
            hasBottomRightNeighbor(in:),
            hasBottomNeighbor(in:)
        ]
    }

    func thElement(in matrix:Matrix) -> Int {
        return matrix[row][col]
    }

    func hasLeftNeighbor(in matrix: Matrix) -> CanMoveValueOfMoveFromPositionFunctionReturnType  {
        guard Position.checkIfIndicesAreValid(in: matrix, row: row, col: col - 1) else {
            return (false, nil)
        }
        return (true, matrix[row][col - 1])
    }

    func hasTopLeftNeighbor(in matrix: Matrix) -> CanMoveValueOfMoveFromPositionFunctionReturnType {
        guard Position.checkIfIndicesAreValid(in: matrix, row: row - 1, col: col - 1) else {
            return (false, nil)
        }
        return (true, matrix[row-1][col - 1])
    }

    func hasTopNeighbor(in matrix: Matrix) -> CanMoveValueOfMoveFromPositionFunctionReturnType {
        guard Position.checkIfIndicesAreValid(in: matrix, row: row - 1, col: col) else {
            return (false, nil)
        }
        return (true, matrix[row-1][col])
    }

    func hasTopRightNeighbor(in matrix: Matrix) -> CanMoveValueOfMoveFromPositionFunctionReturnType {
        guard Position.checkIfIndicesAreValid(in: matrix, row: row - 1, col: col + 1) else {
            return (false, nil)
        }
        return (true, matrix[row-1][col+1])
    }

    func hasRightNeighbor(in matrix: Matrix) -> CanMoveValueOfMoveFromPositionFunctionReturnType {
        guard Position.checkIfIndicesAreValid(in: matrix, row: row, col: col + 1) else {
            return (false, nil)
        }
        return (true, matrix[row][col+1])
    }

    func hasBottomRightNeighbor(in matrix: Matrix) -> CanMoveValueOfMoveFromPositionFunctionReturnType {
        guard Position.checkIfIndicesAreValid(in: matrix, row: row + 1, col: col + 1) else {
            return (false, nil)
        }
        return (true, matrix[row+1][col+1])
    }

    func hasBottomNeighbor(in matrix: Matrix) -> CanMoveValueOfMoveFromPositionFunctionReturnType {
        guard Position.checkIfIndicesAreValid(in: matrix, row: row + 1, col: col) else {
            return (false, nil)
        }
        return (true, matrix[row+1][col])
    }

    func hasBottomLeftNeighbor(in matrix: Matrix) -> CanMoveValueOfMoveFromPositionFunctionReturnType {
        guard Position.checkIfIndicesAreValid(in: matrix, row: row + 1, col: col - 1) else {
            return (false, nil)
        }
        return (true, matrix[row+1][col-1])
    }

    static func checkIfIndicesAreValid(in matrix: Matrix, row: Int, col: Int) -> Bool {
        guard (row >= 0) && (row < matrix.count) else { return false }
        guard (col >= 0) && (col < matrix[row].count) else { return false }
        return true
    }
}

// MARK: - Area
/// Represents an area
struct Area: Equatable {
    /// `startPos` is the starting position of an area e.g. 0, 0
    let startPos: Position
    /// `endPos` is the ending position of an area e.g. 5, 5
    let endPos: Position
}

// MARK: - Result
/// Represents a result object
struct Result: Equatable {
    /// `largestSquareArea` returns the largest square area found in a matrix which is consisting all of defect products.
    let largestSquareArea: Int
    /// `areas` returns the all areas which are square (at least 2x2) and consisting of all of defect products.
    let areas: [Area] // for debugging purposes
    
    static var invalid: Result {
        return Result(largestSquareArea: -1, areas: [])
    }
}

// MARK:  Solution
final class Solution {
    // MARK: - Properties
    private let matrix: Matrix
    
    // MARK: - Lifecycle
    init(matrix: Matrix) {
        self.matrix = matrix
    }
    
    /// Returns the largest square area consisting all of the `ProductType.defect` product type
    func largestSquareDefectArea(matrix: Matrix) -> Result {
        // Empty matrix given
        guard !matrix.isEmpty else { return .invalid }
        
        // Number of rows and columns mismatch in the matrix
        let isMatrixValid: Bool = matrix.reduce(true) { $0 && (matrix.count == $1.count) }
        guard isMatrixValid else {
            return .invalid
        }
        
        var areas: [Area] = []
        let startingPosition = Position(row: 0, col: 0)
        // Starting traversing the matrix from 0, 0
        areas = exploreNearby(matrix: matrix,
                              areas:&areas,
                              offset:0,
                              startingPosition:startingPosition)
        
        // Running one final check to make sure all found areas actually consisting of
        // defect product types
        areas = areas.filter { confirmThat(the: $0, onlyHaving: .defect) }
        
        // Finding the largest square in found list of areas
        let (_, largestAreaLength) = findTheLargestArea(of: areas)
        return Result(largestSquareArea: largestAreaLength, areas: areas)
    }

    /// Iterates over the matrix to find all combinations of a square
    private func exploreNearby(matrix: Matrix, areas: inout [Area], offset:Int, startingPosition:Position) -> [Area] {
        
        // _startingPosition is used internally as a temp variable
        var _startingPosition = startingPosition
        for i in (startingPosition.row)..<(matrix.count - 1) {
            for j in (startingPosition.col)..<(matrix[i].count - 1) {
                
                // Current starting position
                _startingPosition = Position(row: i, col: j)
                
                // Exploring deeper..
                let maxPosition = exploreDeeper(matrix: matrix,
                                                areas: areas,
                                                offset: offset,
                                                startingPosition: _startingPosition)
                
                // If the found maxPosition is greater than the _startingPosition
                // That means we progressed. So we're adding this area to our inout array
                if maxPosition > _startingPosition {
                    areas.append(Area(startPos: _startingPosition, endPos: maxPosition))
                } else {
                    
                }
            }
        }
        
        return areas
    }

    /// This function recursively goes deeper until it does not find defect neighbors
    private func exploreDeeper(matrix:Matrix, areas:[Area], offset:Int, startingPosition:Position) -> Position {
        
        // First checking the immediate neighbors
        let needsToGoDeeper = checkIfImmediateNeighborsAreAlsoDefect(matrix: matrix, curPos: startingPosition)
        
        // If immediate neighbors mismatch, no need to go one level deeper
        guard needsToGoDeeper else {
            return startingPosition
        }
        
        // Checking if we can go bottom-right
        guard startingPosition.hasBottomRightNeighbor(in: matrix).canMove else {
            return startingPosition
        }
        
        let bottomRightPos = Position(row: startingPosition.row+1, col: startingPosition.col+1)
        // Going deeper
        return exploreDeeper(matrix: matrix,
                             areas: areas,
                             offset: offset,
                             startingPosition: bottomRightPos)
    }

    /// Checks if the immediate neighbors (window = 1), are also defect (1)
    /// Example given with i = 1, j = 1:
    /// 1 1 1
    /// 0 1 1
    /// 1 1 1
    /// Compares (1,1) with (0,0), (0,1), (0,2), (1,0), (1,2), (2,0), (2,1) and (2,2) elements and returns true if all of them are defect (1), false otherwise.
    private func checkIfImmediateNeighborsAreAlsoDefect(matrix: Matrix, curPos:Position) -> Bool {
        let currentElement = curPos.thElement(in: matrix)
        // Safety check to make sure current element is actually defect
        guard currentElement == ProductType.defect.rawValue else { return false }
        
        // Iterating over all the neighbors
        // top, top-right, right, right-bottom, bottom, bottom-left, left, top-left
        // checking if we can move there
        // if we can move there, moving there to check the value of it
        return curPos.allMovementsPossible(in: matrix).reduce(true) { (acc, arg1) -> Bool in
            let moveResult = arg1(matrix)
            // If we cannot move we just yield same accumulation
            guard moveResult.canMove else { return acc }
            // If we can move to the neighbor, checking if we have the same value
            return moveResult.valueOfNeighbor == currentElement
        }
    }

    /// This function finds the largest area amongst a list of areas
    private func findTheLargestArea(of areas:[Area]) -> (Area?, Int) {
        var largestAreaLength: Int = -1
        var largestArea: Area? = nil
        for area in areas {
            let lengthOfCurrentArea = length(ofThe: area)
            if lengthOfCurrentArea >= largestAreaLength {
                largestAreaLength = lengthOfCurrentArea
                largestArea = area
            }
        }
        return (largestArea, largestAreaLength)
    }

    /// This function finds the lenght of an area
    /// using Start:(a,b), End:(c,d) => d - b
    private func length(ofThe area:Area) -> Int {
        return area.endPos.col - area.startPos.col + 1
    }

    /// This function runs after all potential areas are found but in some cases there are
    /// underlooked elements and they need to be sent away such as:
    /// [1, 1, 1, 0, 0],
    /// [1, 1, 1, 1, 1],
    /// [1, 1, 1, 1, 1],
    /// [0, 1, 1, 1, 1],
    /// [0, 1, 1, 1, 1],
    /// Currently this algorithm explores diagonally and thinks that this is 5x5 defected square
    /// which it is not
    private func confirmThat(the area:Area, onlyHaving product:ProductType) -> Bool {
        var areaComponents: [Int] = []
        for i in area.startPos.row...area.endPos.row {
            for j in area.startPos.col...area.endPos.col {
                if Position.checkIfIndicesAreValid(in: matrix, row: i, col: j) {
                    areaComponents.append(matrix[i][j])
                }
            }
        }
        return areaComponents.reduce(true) { $0 && $1 == product.rawValue }
    }
}


// MARK: - Test Cases
final class LargestSquareDefectAreaTests: XCTestCase {
    func testInvalidInput() {
        let matrix = [
            [1, 1, 1],
            [1, 1, 0]
        ]
        let s = Solution(matrix: matrix)
        let r = s.largestSquareDefectArea(matrix: matrix)
        XCTAssertEqual(r.largestSquareArea, -1)
        XCTAssertTrue(r.areas.isEmpty)
    }
    
    func testEmptyMatrix() {
        let matrix: Matrix = [[]]
        let s = Solution(matrix: matrix)
        let r = s.largestSquareDefectArea(matrix: matrix)
        XCTAssertEqual(r, Result.invalid)
        XCTAssertTrue(r.areas.isEmpty)
        XCTAssertEqual(r.largestSquareArea, -1)
    }
    
    func test4by4() {
        let matrix = [
            [1, 1, 1, 1],
            [0, 1, 1, 1],
            [0, 1, 1, 1],
            [0, 1, 1, 1]
        ]
        let s = Solution(matrix: matrix)
        let r = s.largestSquareDefectArea(matrix: matrix)
        XCTAssertEqual(r.largestSquareArea, 3)
    }
    
    func test6by6() {
        let matrix = [
            [1, 1, 1, 1, 0, 1],
            [0, 1, 1, 1, 1, 1],
            [0, 1, 1, 1, 1, 1],
            [0, 1, 1, 1, 1, 1],
            [0, 1, 1, 1, 1, 1],
            [0, 1, 1, 1, 0, 1]
        ]
        let s = Solution(matrix: matrix)
        let r = s.largestSquareDefectArea(matrix: matrix)
        XCTAssertEqual(r.largestSquareArea, 4)
    }
}
LargestSquareDefectAreaTests.defaultTestSuite.run()