Magic square diff - Swift

MagicSquare.swift

/*
We define a magic square to be an `n x n` matrix of distinct positive integers from `1` to `n^2` where the sum of any row, column, or diagonal of length `n` is always equal to the same number: the magic constant.

You will be given a `3x3` matrix `s` of integers in the inclusive range `[1,9]`. We can convert any digit  to any other digit `b` in the range `[1,9]` at cost of `|a-b|`. Given `s`, convert it into a magic square at minimal cost. Return this minimal cost.

Note: The resulting magic square must contain distinct integers in the inclusive range `[1,9]`.

Example: $s = [[5, 3, 4], [1, 5, 8], [6, 4, 2]]
The matrix looks like this:
5 3 4
1 5 8
6 4 2

We can convert it to the following magic square:
8 3 4
1 5 9
6 7 2

This took three replacements at a cost of `|5-8|+|8-9|+|4-7|=7`.
*/

func mirror(s: [[Int]]) -> [[Int]] {
    var output = [[Int]]()
    for i in 0..<s.count {
        var row = [Int]()
        for j in 0..<s[i].count {
            row.append(s[i][s[i].count - j - 1])
        }
        output.append(row)
    }
    return output
}

func rotate(s: [[Int]]) -> [[Int]] {
    var output = [[Int]]()
    for j in 0...2 {
        var lst = [Int]()
        for i in 0...2 {
            lst.append(s[2-i][j])
        }
        output.append(lst)
    }
    return output
}

func diff(s: [[Int]], arr: [[Int]]) -> Int {
    var total = 0
    for i in 0...2 {
        for j in 0...2 {
            total += abs(s[i][j] - arr[i][j])            
        }
    }
    return total
}

func formingMagicSquare(s: [[Int]]) -> Int {
    var magic = [[4,3,8], [9,5,1], [2,7,6]]
    var mirrored = mirror(s: magic)
    var totals = [Int]()
    for _ in 0..<4 {
        let diff1 = diff(s: s, arr: magic)
        let diff2 = diff(s: s, arr: mirrored)
        
        totals.append(diff1)
        totals.append(diff2)
        
        magic = rotate(s: magic)
        mirrored = mirror(s: magic)
    }
    return totals.min() ?? 0
}